#!/usr/bin/env python3
"""
Ingest FURNACE design docs into ChromaDB collection 'furnace-design'.

Chunks files at heading boundaries (~1500-2000 chars per chunk).
Adds metadata: source_file, doc_type, component, version, has_code, chunk_index.

Usage:
    # Using the same python env as chroma-mcp:
    /Users/gyan/.cache/uv/archive-v0/wlEdbaN1dC7MI0tn5hJtD/bin/python scripts/ingest_design_docs.py

    # Or with uvx:
    uvx --from chromadb chromadb  # just to get the dep, then:
    python3 scripts/ingest_design_docs.py
"""

import os
import sys
import re
import json
from pathlib import Path

try:
    import chromadb
except ImportError:
    print("ERROR: chromadb not installed. Try running with the chroma-mcp python:")
    print("  /Users/gyan/.cache/uv/archive-v0/wlEdbaN1dC7MI0tn5hJtD/bin/python scripts/ingest_design_docs.py")
    sys.exit(1)

DESIGN_DIR = Path(__file__).parent.parent / "Design"

# File metadata mapping
FILE_METADATA = {
    "FURNACE_Skills_Distribution_Telemetry.md": {
        "doc_type": "guide", "component": "Skills/Telemetry/Distribution", "version": "v3"
    },
    "FURNACE_vs_Claude_Code_Architecture_Analysis.md": {
        "doc_type": "analysis", "component": "Architecture", "version": "v3"
    },
    "FURNACE_Reconciling_Vector_DB_and_Agentic_Search.md": {
        "doc_type": "analysis", "component": "Search", "version": "v4"
    },
    "FURNACE_Parallel_Task_Approval_Workflow.md": {
        "doc_type": "guide", "component": "Parallel_Tasks", "version": "v3"
    },
    "FURNACE_LLD_v4_03_SearchMetrics.md": {
        "doc_type": "lld", "component": "SearchMetrics", "version": "v4"
    },
    "FURNACE_LLD_Part3_Trinity_Legolas.md": {
        "doc_type": "lld", "component": "Trinity/Legolas", "version": "v3"
    },
    "FURNACE_v4_Complete_LLD_Mapping_Guide.md": {
        "doc_type": "guide", "component": "System", "version": "v4"
    },
    "FURNACE_LLD_Part3b_Aragorn_Complete.md": {
        "doc_type": "lld", "component": "Aragorn", "version": "v3"
    },
    "FURNACE_Complete_LLD_v3_with_AUTO-FORGE.md": {
        "doc_type": "lld", "component": "AutoForge", "version": "v3"
    },
    "FURNACE_LLD_v4_01_SearchOrchestrator_Complete.md": {
        "doc_type": "lld", "component": "SearchOrchestrator", "version": "v4"
    },
    "FURNACE_LLD_v4_08_MCP_Server.md": {
        "doc_type": "lld", "component": "MCP_Server", "version": "v4"
    },
    "FURNACE_LLD_v4_04_Gandalf.md": {
        "doc_type": "lld", "component": "Gandalf", "version": "v4"
    },
    "FURNACE_HLD_v4_FINAL_Complete.md": {
        "doc_type": "hld", "component": "System", "version": "v4"
    },
    "FURNACE_LLD_v4_06_Trinity.md": {
        "doc_type": "lld", "component": "Trinity", "version": "v4"
    },
    "FURNACE_LLD_v4_09_Morpheus.md": {
        "doc_type": "lld", "component": "Morpheus", "version": "v4"
    },
    "FURNACE_LLD_v4_05_Gimli.md": {
        "doc_type": "lld", "component": "Gimli", "version": "v4"
    },
    "FURNACE_LLD_v4_07_Elrond_FreshnessManager.md": {
        "doc_type": "lld", "component": "Elrond/FreshnessManager", "version": "v4"
    },
    "FURNACE_LLD_v4_02_AgenticSearch_Detailed.md": {
        "doc_type": "lld", "component": "AgenticSearch", "version": "v4"
    },
}


def slugify(filename: str) -> str:
    """Convert filename to a URL-safe slug for use as chunk ID prefix."""
    name = filename.replace(".md", "")
    name = re.sub(r"[^a-zA-Z0-9]+", "-", name)
    name = name.strip("-").lower()
    return name


def chunk_markdown(content: str, target_size: int = 1800, min_size: int = 200) -> list[str]:
    """Split markdown content at heading boundaries into chunks of ~target_size chars.

    Strategy:
    1. Split at ## and ### headings
    2. Merge small chunks with the next one
    3. Split overly large chunks at paragraph boundaries
    """
    # Split at heading boundaries (## or ###)
    sections = re.split(r'(?=^#{2,3}\s)', content, flags=re.MULTILINE)

    chunks = []
    current_chunk = ""

    for section in sections:
        section = section.strip()
        if not section:
            continue

        # If adding this section keeps us under target, accumulate
        if len(current_chunk) + len(section) < target_size:
            if current_chunk:
                current_chunk += "\n\n" + section
            else:
                current_chunk = section
        else:
            # Save current chunk if it's big enough
            if len(current_chunk) >= min_size:
                chunks.append(current_chunk.strip())
                current_chunk = section
            else:
                # Current chunk too small, merge with section
                if current_chunk:
                    current_chunk += "\n\n" + section
                else:
                    current_chunk = section

            # If accumulated chunk is too large, split at paragraph boundaries
            while len(current_chunk) > target_size * 1.5:
                # Find a good split point (double newline near target)
                split_pos = current_chunk.rfind("\n\n", 0, target_size)
                if split_pos == -1:
                    split_pos = current_chunk.rfind("\n", 0, target_size)
                if split_pos == -1:
                    split_pos = target_size

                chunks.append(current_chunk[:split_pos].strip())
                current_chunk = current_chunk[split_pos:].strip()

    # Don't forget the last chunk
    if current_chunk.strip():
        # If it's very small, merge with the previous chunk
        if len(current_chunk) < min_size and chunks:
            chunks[-1] += "\n\n" + current_chunk.strip()
        else:
            chunks.append(current_chunk.strip())

    return chunks


def has_code_block(text: str) -> bool:
    """Check if text contains code blocks."""
    return "```" in text


def find_chromadb_path() -> str:
    """Try to find where the chroma-mcp server stores its data."""
    # Common locations to check
    candidates = [
        os.path.expanduser("~/chroma_data"),
        os.path.expanduser("~/.chroma"),
        "/tmp/chroma",
        os.path.expanduser("~/Library/Application Support/chroma"),
    ]

    for path in candidates:
        if os.path.exists(path) and os.path.isdir(path):
            return path

    return None


def main():
    print(f"ChromaDB version: {chromadb.__version__}")
    print(f"Design dir: {DESIGN_DIR}")
    print(f"Files to process: {len(FILE_METADATA)}")
    print()

    # Try to connect to the same ChromaDB as chroma-mcp
    # The chroma-mcp package uses HttpClient by default connecting to localhost:8000
    # But if run as PersistentClient, it stores in current directory

    # First try HTTP client (if chroma server is running)
    client = None

    try:
        client = chromadb.HttpClient(host="localhost", port=8000)
        # Test connection
        client.heartbeat()
        print("Connected to ChromaDB via HTTP (localhost:8000)")
    except Exception as e:
        print(f"HTTP connection failed: {e}")
        print("Trying PersistentClient...")

        # Try PersistentClient at various locations
        chroma_path = find_chromadb_path()
        if chroma_path:
            print(f"Found ChromaDB data at: {chroma_path}")
            client = chromadb.PersistentClient(path=chroma_path)
        else:
            # Use default (current directory)
            print("Using default PersistentClient (./chroma_data)")
            client = chromadb.PersistentClient()

    # List existing collections
    collections = client.list_collections()
    print(f"Existing collections: {[c.name for c in collections]}")

    # Get or create the furnace-design collection
    try:
        collection = client.get_collection("furnace-design")
        existing_count = collection.count()
        print(f"Found existing collection 'furnace-design' with {existing_count} documents")

        # Delete all existing documents to do a fresh ingest
        if existing_count > 0:
            print(f"Clearing {existing_count} existing documents...")
            # Get all IDs and delete them
            existing = collection.get(limit=existing_count)
            if existing["ids"]:
                collection.delete(ids=existing["ids"])
                print(f"Deleted {len(existing['ids'])} documents")
    except Exception:
        print("Creating new collection 'furnace-design'...")
        collection = client.create_collection("furnace-design")

    print()

    # Process each file
    total_chunks = 0
    batch_docs = []
    batch_ids = []
    batch_metadatas = []
    BATCH_SIZE = 20

    for filename, meta in FILE_METADATA.items():
        filepath = DESIGN_DIR / filename

        if not filepath.exists():
            print(f"WARNING: File not found: {filepath}")
            continue

        content = filepath.read_text(encoding="utf-8")
        slug = slugify(filename)

        # Chunk the content
        chunks = chunk_markdown(content)

        print(f"  {filename}: {len(chunks)} chunks ({len(content)} chars)")

        for i, chunk_text in enumerate(chunks):
            chunk_id = f"{slug}-chunk-{i}"

            metadata = {
                "source_file": filename,
                "doc_type": meta["doc_type"],
                "component": meta["component"],
                "version": meta["version"],
                "has_code": has_code_block(chunk_text),
                "chunk_index": i,
                "total_chunks": len(chunks),
                "char_count": len(chunk_text),
            }

            batch_docs.append(chunk_text)
            batch_ids.append(chunk_id)
            batch_metadatas.append(metadata)
            total_chunks += 1

            # Flush batch when full
            if len(batch_docs) >= BATCH_SIZE:
                collection.add(
                    documents=batch_docs,
                    ids=batch_ids,
                    metadatas=batch_metadatas,
                )
                print(f"    Ingested batch of {len(batch_docs)} chunks (total: {total_chunks})")
                batch_docs = []
                batch_ids = []
                batch_metadatas = []

    # Flush remaining batch
    if batch_docs:
        collection.add(
            documents=batch_docs,
            ids=batch_ids,
            metadatas=batch_metadatas,
        )
        print(f"    Ingested final batch of {len(batch_docs)} chunks (total: {total_chunks})")

    print()
    print(f"DONE! Total chunks ingested: {total_chunks}")
    print(f"Collection count: {collection.count()}")

    # Quick verification query
    print()
    print("Verification query: 'SearchOrchestrator tiered search'")
    results = collection.query(
        query_texts=["SearchOrchestrator tiered search"],
        n_results=3,
    )

    for i, (doc_id, meta) in enumerate(zip(results["ids"][0], results["metadatas"][0])):
        print(f"  {i+1}. {doc_id} ({meta['source_file']}, {meta['component']})")


if __name__ == "__main__":
    main()
