#!/usr/bin/env python3
"""Walk the intranet/ folder and ingest every text-bearing file into AnythingLLM.

Pipeline per file:
  1. POST /api/v1/document/upload   -> AnythingLLM parses + chunks + saves
                                        metadata, returns documents:[{location}]
  2. Collect all returned locations.
  3. POST /api/v1/workspace/intranet/update-embeddings with adds=locations
     -> AnythingLLM embeds each chunk via Ollama bge-m3 and indexes into LanceDB.

We do step 1 sequentially (server processes serially anyway), then a single
bulk step 2 so the heavy embedding pass is one operation we can watch.

Skipped extensions: jpg/png (no OCR), mp4 (no whisper for now), crt, vsdx.
"""
import os
import sys
import time
from pathlib import Path

import requests

BASE = os.environ.get("ANYTHINGLLM_URL", "http://127.0.0.1:3035")
API_KEY = os.environ.get("API_KEY",
    "AILM-78c5154c6dc4fc04a940784c8ad9927e8882c7e16bf4cdcd044ccd0fd00af168")
WORKSPACE = os.environ.get("WORKSPACE", "intranet")
ROOT = Path(os.environ.get("INTRANET_ROOT",
    "/home/klas/Prace/AI/portal/anythingllm/intranet/INTRANET"))

ALLOWED_EXTS = {".pdf", ".docx", ".doc", ".xlsx", ".xls",
                ".pptx", ".txt", ".html", ".md"}

HEADERS = {"Authorization": f"Bearer {API_KEY}"}


def upload_file(path: Path) -> str | None:
    """Upload one file. Returns its document 'location' on success."""
    with open(path, "rb") as f:
        files = {"file": (path.name, f, "application/octet-stream")}
        r = requests.post(
            f"{BASE}/api/v1/document/upload",
            headers=HEADERS,
            files=files,
            data={"addToWorkspaces": ""},  # we'll attach in bulk after
            timeout=300,
        )
    if not r.ok:
        print(f"  ! HTTP {r.status_code}: {r.text[:200]}")
        return None
    data = r.json()
    if not data.get("success"):
        print(f"  ! upload failed: {data.get('error', data)}")
        return None
    docs = data.get("documents") or []
    if not docs:
        print("  ! no documents in response")
        return None
    return docs[0]["location"]


def attach_to_workspace(locations: list[str]) -> dict:
    """Embed + index everything we uploaded into the workspace."""
    r = requests.post(
        f"{BASE}/api/v1/workspace/{WORKSPACE}/update-embeddings",
        headers={**HEADERS, "Content-Type": "application/json"},
        json={"adds": locations, "deletes": []},
        timeout=3600,
    )
    r.raise_for_status()
    return r.json()


def main() -> int:
    if not ROOT.exists():
        print(f"Root does not exist: {ROOT}", file=sys.stderr)
        return 1

    files = sorted(
        p for p in ROOT.rglob("*")
        if p.is_file() and p.suffix.lower() in ALLOWED_EXTS
    )
    print(f"Found {len(files)} ingestible files under {ROOT}")

    locations: list[str] = []
    failed: list[Path] = []
    t0 = time.time()
    for i, path in enumerate(files, 1):
        rel = path.relative_to(ROOT)
        size_kb = path.stat().st_size / 1024
        print(f"[{i:>3}/{len(files)}] {rel}  ({size_kb:.0f} KB)")
        loc = upload_file(path)
        if loc:
            locations.append(loc)
        else:
            failed.append(path)

    print(f"\nUploaded {len(locations)}/{len(files)} files in "
          f"{time.time()-t0:.0f}s. Failed: {len(failed)}")
    if failed:
        print("Failed files:")
        for p in failed:
            print(f"  - {p.relative_to(ROOT)}")

    if not locations:
        print("Nothing to embed.")
        return 1

    print(f"\nEmbedding {len(locations)} docs into workspace '{WORKSPACE}'...")
    t0 = time.time()
    result = attach_to_workspace(locations)
    print(f"Done in {time.time()-t0:.0f}s")
    print(f"Workspace now has {len(result.get('workspace', {}).get('documents', []))} documents")
    return 0


if __name__ == "__main__":
    sys.exit(main())