#!/usr/bin/env python3
"""
vault_rag_expand.py — Atlas Mind expansion.

Walks selected vault folders + batch uploads .md files to Open WebUI KB
so ALL cluster models query against the same knowledge graph.

Curated set (signal > noise):
  - 00_Atlas_Wiki/*    — operator-curated canonical pages
  - 01_Clients/*       — client briefs + relationship history
  - 04_Sales/*         — pipeline, pricing, templates
  - 06_Marketing/*     — brand, positioning, campaigns
  - 07_Finance/*.md    — money facts + decisions

Skip:
  - 99_Quarantine/ + 99_Archive/ (dead)
  - .git/ .obsidian/ (config)
  - Files > 100 KB (template-heavy, low signal)
"""
import os
import json
import time
import hashlib
import urllib.request
import urllib.error
from pathlib import Path

OPENWEBUI = "http://<YOUR_VPS_IP>:8082"
KB_ID = "c2b9f985-69f6-4dee-b216-73d15bdfb8bd"
TOKEN_PATH = Path(r"C:\Users\Gebruiker\AppData\Local\Atlas\openwebui-token.json")
VAULT_ROOT = Path(r"C:\Users\Gebruiker\MEGA\Atlas Corporation")
INDEX_LOG = Path(r"C:\Users\Gebruiker\AppData\Local\Atlas\analyses\rag_index.json")

INCLUDE_FOLDERS = [
    "00_Atlas_Wiki",
    "01_Clients",
    "04_Sales",
    "06_Marketing",
    "07_Finance",
]
EXCLUDE_KEYWORDS = ["99_Quarantine", "99_Archive", "_archive", ".git", ".obsidian",
                    "backup", "_backup", "draft-202", "old_"]
MAX_FILE_SIZE = 100_000  # 100 KB
MAX_FILES_PER_RUN = 200  # rate-limit per cron tick


def load_token():
    with open(TOKEN_PATH) as f:
        return json.load(f)["token"]


def load_index():
    if INDEX_LOG.exists():
        try: return json.loads(INDEX_LOG.read_text(encoding="utf-8"))
        except: pass
    return {"uploaded": {}, "errors": []}


def save_index(idx):
    INDEX_LOG.parent.mkdir(parents=True, exist_ok=True)
    INDEX_LOG.write_text(json.dumps(idx, indent=2), encoding="utf-8")


def file_hash(p):
    return hashlib.sha1(p.read_bytes()).hexdigest()[:12]


def list_candidates():
    """Walk include-folders, yield (Path, relpath_str, hash)."""
    out = []
    for folder in INCLUDE_FOLDERS:
        root = VAULT_ROOT / folder
        if not root.exists():
            print(f"skip (not found): {root}")
            continue
        for p in root.rglob("*.md"):
            rel = p.relative_to(VAULT_ROOT).as_posix()
            if any(kw in rel for kw in EXCLUDE_KEYWORDS): continue
            if p.stat().st_size > MAX_FILE_SIZE: continue
            if p.stat().st_size < 100: continue  # near-empty
            try:
                h = file_hash(p)
                out.append((p, rel, h))
            except Exception:
                pass
    return out


def upload_file(token, path: Path, relpath: str):
    """Upload single file to Open WebUI via multipart, then attach to KB."""
    # Build multipart manually (urllib doesn't have it natively, use boundary)
    import uuid
    boundary = "----atlas-" + uuid.uuid4().hex
    filename = path.name
    file_bytes = path.read_bytes()
    body_parts = [
        f"--{boundary}".encode(),
        f'Content-Disposition: form-data; name="file"; filename="{filename}"'.encode(),
        b"Content-Type: text/markdown",
        b"",
        file_bytes,
        f"--{boundary}--".encode(),
        b"",
    ]
    body = b"\r\n".join(body_parts)

    # Step 1: upload file
    req = urllib.request.Request(
        f"{OPENWEBUI}/api/v1/files/",
        data=body, method="POST",
    )
    req.add_header("Authorization", f"Bearer {token}")
    req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            file_info = json.loads(r.read().decode())
        file_id = file_info.get("id")
        if not file_id:
            return False, "no file_id returned"
    except urllib.error.HTTPError as e:
        return False, f"upload HTTP {e.code}: {e.read().decode()[:120]}"
    except Exception as e:
        return False, f"upload err: {type(e).__name__}"

    # Step 2: attach to KB
    req2 = urllib.request.Request(
        f"{OPENWEBUI}/api/v1/knowledge/{KB_ID}/file/add",
        data=json.dumps({"file_id": file_id}).encode(),
        method="POST",
    )
    req2.add_header("Authorization", f"Bearer {token}")
    req2.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(req2, timeout=30) as r:
            r.read()
        return True, file_id
    except urllib.error.HTTPError as e:
        return False, f"attach HTTP {e.code}: {e.read().decode()[:120]}"
    except Exception as e:
        return False, f"attach err: {type(e).__name__}"


def main():
    token = load_token()
    idx = load_index()
    cands = list_candidates()
    print(f"Found {len(cands)} candidate files in {len(INCLUDE_FOLDERS)} folders")

    new_to_upload = []
    skipped = 0
    for p, rel, h in cands:
        prev = idx["uploaded"].get(rel)
        if prev and prev.get("hash") == h:
            skipped += 1
            continue
        new_to_upload.append((p, rel, h))

    new_to_upload = new_to_upload[:MAX_FILES_PER_RUN]
    print(f"Already uploaded (skip): {skipped}")
    print(f"To upload this run: {len(new_to_upload)} (cap {MAX_FILES_PER_RUN})")

    ok_count = 0
    err_count = 0
    t_start = time.time()
    for i, (p, rel, h) in enumerate(new_to_upload, 1):
        success, info = upload_file(token, p, rel)
        if success:
            idx["uploaded"][rel] = {"file_id": info, "hash": h, "ts": time.strftime("%Y-%m-%d %H:%M:%S")}
            ok_count += 1
            print(f"  [{i}/{len(new_to_upload)}] OK {rel}")
        else:
            idx["errors"].append({"rel": rel, "err": str(info), "ts": time.strftime("%Y-%m-%d %H:%M:%S")})
            err_count += 1
            print(f"  [{i}/{len(new_to_upload)}] FAIL {rel} -> {info}")
        # Save state every 10 uploads in case we crash
        if i % 10 == 0:
            save_index(idx)
        time.sleep(0.3)  # gentle pacing

    save_index(idx)
    elapsed = round(time.time() - t_start, 1)
    print(f"\nDone in {elapsed}s")
    print(f"  uploaded: {ok_count}")
    print(f"  errors: {err_count}")
    print(f"  total in index: {len(idx['uploaded'])}")


if __name__ == "__main__":
    main()