"""
Codara Real Benchmark
Compares naive file loading vs Codara context engine
on the jcodemunch repository.
"""

import os
import sys
import time
from pathlib import Path

# Add codara to path
sys.path.insert(0, str(Path(__file__).parent))

from kodara.indexer.scanner import RepositoryScanner, LANGUAGE_EXTENSIONS, SKIP_DIRS
from kodara.memory.store import MemoryStore
from kodara.context_engine.engine import ContextEngine

REPO_PATH = Path("D:/Desktop/Codara/jcodmach-main/jcodmach-main")
REPO_ID = "local/jcodmach-main-45b013cb"

QUERIES = [
    "How does symbol extraction and AST parsing work?",
    "How does storage and indexing work?",
    "How does search and retrieval work?",
    "How does security and path traversal protection work?",
    "How does AI summarization work?",
]


def count_tokens(text) -> int:
    """Rough token estimate: 4 chars per token."""
    if isinstance(text, int):
        return text // 4
    return len(text) // 4


def naive_load_all_files(repo_path: Path) -> dict:
    """
    Simulate naive approach: load ALL source files into context.
    This is what AI tools do without a memory layer.
    """
    total_chars = 0
    total_files = 0
    file_sizes = []

    for path in repo_path.rglob("*"):
        if not path.is_file():
            continue

        # Skip same dirs as our scanner
        rel_parts = path.relative_to(repo_path).parts
        if any(p in SKIP_DIRS or p.startswith(".") for p in rel_parts[:-1]):
            continue

        ext = path.suffix.lower()
        if ext not in LANGUAGE_EXTENSIONS:
            continue

        try:
            content = path.read_text(encoding="utf-8", errors="replace")
            size = len(content)
            total_chars += size
            total_files += 1
            file_sizes.append((path.name, size))
        except Exception:
            continue

    file_sizes.sort(key=lambda x: -x[1])

    return {
        "total_files": total_files,
        "total_chars": total_chars,
        "total_tokens": count_tokens(total_chars),
        "top_files": file_sizes[:5],
    }


def codara_context_for_query(query: str, engine, store) -> dict:
    """Run a single Codara query and measure output size."""
    t0 = time.time()
    result = engine.query(REPO_ID, query, max_modules=8, depth=1)
    formatted = engine.format_context(result)
    elapsed = time.time() - t0

    return {
        "query": query,
        "modules_selected": len(result.modules),
        "total_modules": result.total_modules_searched,
        "output_chars": len(formatted),
        "output_tokens": count_tokens(formatted),
        "elapsed_ms": round(elapsed * 1000, 1),
    }


def main():
    print("=" * 60)
    print("CODARA REAL BENCHMARK")
    print("=" * 60)
    print(f"Repository: {REPO_PATH.name}")
    print()

    # ── 1. Naive approach ──────────────────────────────────────
    print("Measuring naive approach (load all files)...")
    naive = naive_load_all_files(REPO_PATH)

    print("\n" + "-"*60)
    print("NAIVE APPROACH (no memory layer)")
    print("-"*60)
    print(f"  Files loaded:   {naive['total_files']}")
    print(f"  Total chars:    {naive['total_chars']:,}")
    print(f"  Total tokens:   {naive['total_tokens']:,}")
    print(f"\n  Top 5 largest files:")
    for name, size in naive["top_files"]:
        print(f"    {name:<40} {size:>8,} chars  (~{count_tokens(size):,} tokens)")

    # ── 2. Codara approach ─────────────────────────────────────
    store = MemoryStore()
    engine = ContextEngine(store)

    # Check if indexed
    idx = store.load(REPO_ID)
    if not idx:
        print("\nIndexing repository first...")
        from kodara.indexer.scanner import RepositoryScanner
        from kodara.graph.dependency import DependencyGraph
        from kodara.memory.store import KodaraIndex
        from kodara.summarizer import generate_summaries

        scanner = RepositoryScanner(REPO_PATH)
        scan = scanner.scan()
        graph = DependencyGraph()
        graph.build(scan)
        graph_data = graph.to_dict()

        modules = {}
        for fp, info in scan.files.items():
            modules[fp] = {
                "path": fp,
                "module_name": graph.node_data(fp).get("module_name", ""),
                "language": info.language,
                "classes": info.classes,
                "functions": info.functions,
                "exports": info.exports,
                "imports": info.imports,
                "file_hash": info.hash,
                "summary": "",
                "dependencies": graph.get_dependencies(fp),
                "dependents": graph.get_dependents(fp),
            }

        idx = KodaraIndex(
            repo_id=REPO_ID,
            display_name=REPO_PATH.name,
            source_root=str(REPO_PATH),
            indexed_at=time.time(),
            index_version=1,
            file_count=len(scan.files),
            language_counts=scan.languages,
            modules=modules,
            graph_edges=graph_data["edges"],
            file_hashes={p: f.hash for p, f in scan.files.items()},
        )
        store.save(idx)
        print("Done.\n")

    print("\n" + "-"*60)
    print("CODARA APPROACH (architecture memory)")
    print("-"*60)
    print(f"  Total indexed modules: {len(idx.modules)}")
    print()

    results = []
    for query in QUERIES:
        r = codara_context_for_query(query, engine, store)
        results.append(r)

        reduction_pct = 100 * (1 - r["output_tokens"] / naive["total_tokens"])
        ratio = naive["total_tokens"] / max(r["output_tokens"], 1)

        print(f"  Query: \"{query[:50]}\"")
        print(f"    Modules:  {r['modules_selected']}/{r['total_modules']} selected")
        print(f"    Tokens:   {r['output_tokens']:,}  (vs {naive['total_tokens']:,} naive)")
        print(f"    Saving:   {reduction_pct:.1f}%  |  {ratio:.1f}x fewer tokens")
        print(f"    Time:     {r['elapsed_ms']}ms")
        print()

    # ── 3. Summary ─────────────────────────────────────────────
    avg_codara_tokens = sum(r["output_tokens"] for r in results) / len(results)
    avg_reduction = 100 * (1 - avg_codara_tokens / naive["total_tokens"])
    avg_ratio = naive["total_tokens"] / max(avg_codara_tokens, 1)

    print("="*60)
    print("SUMMARY")
    print("="*60)
    print(f"  Naive (all files):     {naive['total_tokens']:>8,} tokens")
    print(f"  Codara avg per query:  {avg_codara_tokens:>8,.0f} tokens")
    print(f"  Average reduction:     {avg_reduction:>7.1f}%")
    print(f"  Average ratio:         {avg_ratio:>7.1f}x fewer tokens")
    print()
    print(f"  Result: Codara delivers ~{avg_ratio:.0f}x token reduction")
    print(f"          on the {REPO_PATH.name} repository.")
    print("=" * 60)

    return {
        "naive_tokens": naive["total_tokens"],
        "codara_avg_tokens": avg_codara_tokens,
        "avg_reduction_pct": avg_reduction,
        "avg_ratio": avg_ratio,
        "queries": results,
    }


if __name__ == "__main__":
    main()
