"""CLI entry point for the RAG subsystem.

Usage::

    python -m axiom.rag index [path]        # index workspace docs
    python -m axiom.rag search "query"      # hybrid search
    python -m axiom.rag status              # per-corpus statistics
    python -m axiom.rag load-community      # load community corpus dump
    python -m axiom.rag sync org            # sync org corpus
    python -m axiom.rag reindex             # force re-index all docs

Legacy aliases: ``ingest`` → ``index``, ``stats`` → ``status``
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

log = logging.getLogger(__name__)


def _load_env() -> None:
    """Try to load DATABASE_URL from .env or axi settings if not already set."""
    if os.environ.get("DATABASE_URL"):
        return
    # Try .env file
    for candidate in (Path.cwd() / ".env", Path(__file__).resolve().parents[4] / ".env"):
        if candidate.is_file():
            for line in candidate.read_text().splitlines():
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, _, val = line.partition("=")
                    key, val = key.strip(), val.strip().strip("\"'")
                    if key and key not in os.environ:
                        os.environ[key] = val
            break
    # Fall back to axi settings
    if not os.environ.get("DATABASE_URL"):
        try:
            from axiom.extensions.builtins.settings.store import SettingsStore

            url = SettingsStore().get("rag.database_url", "")
            if url:
                os.environ["DATABASE_URL"] = url
        except Exception:
            pass

    # Fall back to auto-provisioned PG with stored password
    if not os.environ.get("DATABASE_URL"):
        try:
            from axiom.setup.secrets import get_secret

            pg_pass = get_secret("AXIOM_PG_PASSWORD")
            if pg_pass:
                os.environ["DATABASE_URL"] = f"postgresql://axiom:{pg_pass}@localhost:5432/axiom_db"
        except Exception:
            pass


def _get_store():
    from .store import RAGStore

    url = os.environ.get("DATABASE_URL")
    if not url:
        print(
            "ERROR: No RAG database configured.\n"
            '  Set via: axi settings set rag.database_url "postgresql://..."\n'
            "  Or:      export DATABASE_URL=postgresql://...",
            file=sys.stderr,
        )
        sys.exit(1)
    store = RAGStore(url)
    store.connect()
    return store


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------


def cmd_index(args: argparse.Namespace) -> None:
    """Index one or more paths (or the default repo paths)."""
    from axiom import REPO_ROOT

    from .ingest import ingest_file, ingest_repo
    from .store import CORPUS_INTERNAL

    corpus = getattr(args, "corpus", CORPUS_INTERNAL)
    store = _get_store()
    try:
        if args.paths:
            from .extract import SUPPORTED_EXTENSIONS
            from .ingest import IngestStats, ingest_file

            total = IngestStats()
            for raw in args.paths:
                p = Path(raw).resolve()
                if p.is_dir():
                    files = []
                    for ext in SUPPORTED_EXTENSIONS:
                        files.extend(p.rglob(f"*{ext}"))
                    files = [
                        f for f in files if "__MACOSX" not in str(f) and not f.name.startswith(".")
                    ]
                    log.info("Found %d files under %s", len(files), p)
                    for fpath in sorted(files):
                        s = ingest_file(fpath, store, repo_root=p, corpus=corpus)
                        total.files_indexed += s.files_indexed
                        total.chunks_created += s.chunks_created
                        total.files_skipped += s.files_skipped
                elif p.is_file():
                    s = ingest_file(p, store, repo_root=p.parent, corpus=corpus)
                    total.files_indexed += s.files_indexed
                    total.chunks_created += s.chunks_created
                    total.files_skipped += s.files_skipped
                else:
                    print(f"WARNING: path not found: {p}", file=sys.stderr)
        else:
            total = ingest_repo(REPO_ROOT, store, corpus=corpus)

        print(
            f"Indexed: {total.files_indexed} files, "
            f"{total.chunks_created} chunks  "
            f"({total.files_skipped} unchanged/skipped)  "
            f"[corpus: {corpus}]"
        )
    finally:
        store.close()


def cmd_search(args: argparse.Namespace) -> None:
    from .embeddings import embed_texts

    store = _get_store()
    try:
        try:
            embs = embed_texts([args.query])
        except Exception as exc:
            log.warning("Embedding skipped (using text-only search): %s", exc)
            embs = None
        results = store.search(
            query_embedding=embs[0] if embs else None,
            query_text=args.query,
            limit=args.limit,
        )
        if not results:
            print("No results.")
            return

        for i, r in enumerate(results, 1):
            print(f"\n--- Result {i} (score: {r.combined_score:.4f}, corpus: {r.corpus}) ---")
            print(f"  Source: {r.source_path} (chunk {r.chunk_index})")
            print(f"  Title:  {r.source_title}")
            print(f"  {r.chunk_text[:300]}...")
    finally:
        store.close()


def cmd_status(args: argparse.Namespace) -> None:
    """Show per-corpus statistics."""
    store = _get_store()
    try:
        s = store.stats()
        print("RAG Index Status")
        print(f"  Total documents: {s['total_documents']}")
        print(f"  Total chunks:    {s['total_chunks']}")
        print()
        print(f"  {'Corpus':<20} {'Docs':>6} {'Chunks':>8}")
        print(f"  {'-' * 20} {'-' * 6} {'-' * 8}")
        for corpus in ("rag-community", "rag-org", "rag-internal"):
            docs = s["documents_by_corpus"].get(corpus, 0)
            chunks = s["chunks_by_corpus"].get(corpus, 0)
            label = corpus.replace("rag-", "")
            print(f"  {label:<20} {docs:>6} {chunks:>8}")
    finally:
        store.close()


def cmd_load_community(args: argparse.Namespace) -> None:
    """Load the community corpus from a pre-built dump file."""
    from pathlib import Path

    store = _get_store()
    try:
        # If no path given, look for bundled dump
        if args.dump_path:
            dump = Path(args.dump_path)
        else:
            # Look for bundled community dump
            try:
                pkg_dir = Path(__file__).resolve().parents[1] / "data" / "rag"
                candidates = sorted(pkg_dir.glob("community-v*.sql")) + sorted(
                    pkg_dir.glob("community-v*.pgdump")
                )
                if not candidates:
                    print(
                        "No community corpus dump found.\n"
                        "  Expected: src/neutron_os/data/rag/community-v*.sql\n"
                        "  Download: axi update --community-rag",
                        file=sys.stderr,
                    )
                    sys.exit(1)
                dump = candidates[-1]  # latest by name sort
            except Exception as exc:
                print(f"ERROR locating community dump: {exc}", file=sys.stderr)
                sys.exit(1)

        print(f"Loading community corpus from {dump} ...")
        store.load_community_dump(dump)
        s = store.stats()
        chunks = s["chunks_by_corpus"].get("rag-community", 0)
        docs = s["documents_by_corpus"].get("rag-community", 0)
        print(f"Community corpus loaded: {docs} documents, {chunks} chunks")
    finally:
        store.close()


def cmd_sync(args: argparse.Namespace) -> None:
    """Sync org corpus (v1: manual rsync instructions)."""
    target = getattr(args, "target", "org")
    if target == "org":
        # Configure your org corpus source in runtime/config/rag.toml
        # Example: org_corpus_source = "server.example.com:/shared/org-corpus/"
        print(
            "To sync the org corpus, run:\n"
            "\n"
            "  rsync -avz --progress \\\n"
            "    <your-server>:/path/to/org-corpus/ \\\n"
            "    runtime/knowledge/org-corpus/\n"
            "\n"
            "  axi rag index runtime/knowledge/org-corpus/ --corpus rag-org\n"
            "\n"
            "Configure your server path in runtime/config/rag.toml.\n"
            "Automated sync (axi rag sync org --auto) will be available in v2."
        )
    else:
        print(f"Unknown sync target: {target}", file=sys.stderr)
        sys.exit(1)


def cmd_watch(args: argparse.Namespace) -> None:
    """Watch workspace directories and re-index changed files automatically."""
    from axiom import REPO_ROOT

    from .store import CORPUS_INTERNAL
    from .watcher import watch

    corpus = getattr(args, "corpus", CORPUS_INTERNAL)
    store = _get_store()
    try:
        watch(REPO_ROOT, store, corpus=corpus, quiet=args.quiet)
    finally:
        store.close()


def cmd_reindex(args: argparse.Namespace) -> None:
    """Force re-index by clearing checksums and re-running index."""
    from axiom import REPO_ROOT

    from .ingest import ingest_repo
    from .store import CORPUS_INTERNAL

    corpus = getattr(args, "corpus", CORPUS_INTERNAL)
    store = _get_store()
    try:
        # Clear existing corpus so checksums don't skip anything
        deleted = store.delete_corpus(corpus)
        print(f"Cleared {deleted} chunks from {corpus}")
        stats = ingest_repo(REPO_ROOT, store, corpus=corpus)
        print(
            f"Re-indexed: {stats.files_indexed} files, "
            f"{stats.chunks_created} chunks  [corpus: {corpus}]"
        )
    finally:
        store.close()


# Legacy alias
def cmd_ingest(args: argparse.Namespace) -> None:
    """Legacy alias for cmd_index."""
    if not hasattr(args, "paths"):
        args.paths = []
    cmd_index(args)


# ---------------------------------------------------------------------------
# Generation management commands
# ---------------------------------------------------------------------------


def cmd_upgrade(args: argparse.Namespace) -> None:
    """Build a new RAG generation with semantic chunking."""
    from pathlib import Path

    from .upgrade import build_generation

    store = _get_store()
    try:
        stats = build_generation(
            store=store,
            source_path=Path(args.source),
            corpus=args.corpus,
            chunking_tier=args.tier,
            auto_backup=not args.no_backup,
        )
        if stats.success:
            print(
                f"Generation {stats.generation} built: {stats.files_processed} files, "
                f"{stats.chunks_created} chunks (tier={stats.chunking_tier})"
            )
            if stats.backup_path:
                print(f"  Backup: {stats.backup_path}")
        else:
            print(f"Build failed: {stats.error}", file=sys.stderr)
            sys.exit(1)
    finally:
        store.close()


def cmd_promote(args: argparse.Namespace) -> None:
    store = _get_store()
    try:
        from .generation import GenerationManager

        gm = GenerationManager(store)
        gm.promote(args.corpus, args.generation)
        print(f"Promoted generation {args.generation} for {args.corpus}")
    finally:
        store.close()


def cmd_rollback(args: argparse.Namespace) -> None:
    store = _get_store()
    try:
        from .generation import GenerationManager

        gm = GenerationManager(store)
        gm.rollback(args.corpus, args.generation)
        print(f"Rolled back {args.corpus} to generation {args.generation}")
    finally:
        store.close()


def cmd_discard(args: argparse.Namespace) -> None:
    store = _get_store()
    try:
        from .generation import GenerationManager

        gm = GenerationManager(store)
        gm.discard(args.corpus, args.generation)
        print(f"Discarded generation {args.generation} for {args.corpus}")
    finally:
        store.close()


def cmd_benchmark(args: argparse.Namespace) -> None:
    """Run A/B quality benchmark."""
    store = _get_store()
    try:
        from .benchmark import run_ab_benchmark

        # Use a minimal set of gold queries for quick evaluation
        gold = [
            {"query": "emergency core cooling systems", "keywords": ["ECCS", "cooling", "core"]},
            {"query": "MSRE fuel salt composition", "keywords": ["LiF", "BeF2", "fuel", "salt"]},
            {"query": "TRIGA fuel element", "keywords": ["zirconium", "hydride", "TRIGA"]},
            {"query": "radiation protection dose limits", "keywords": ["ALARA", "dose", "rem"]},
            {"query": "criticality safety validation", "keywords": ["keff", "subcritical"]},
            {"query": "neutron transport methods", "keywords": ["neutron", "transport"]},
            {"query": "reactor safety analysis", "keywords": ["safety", "analysis", "reactor"]},
            {"query": "molten salt reactor", "keywords": ["molten", "salt", "MSR"]},
            {"query": "graphite moderator", "keywords": ["graphite", "moderator"]},
            {"query": "NRC regulatory requirements", "keywords": ["NRC", "regulation"]},
        ]
        tiers = args.compare
        report = run_ab_benchmark(store, gold, tier_a=tiers[0], tier_b=tiers[1])
        print(report.summary)
    finally:
        store.close()


def cmd_generations(args: argparse.Namespace) -> None:
    """Show generation status per corpus."""
    from .store import ALL_CORPORA

    store = _get_store()
    try:
        from .generation import GenerationManager

        gm = GenerationManager(store)
        print("Corpus               Active  Candidate")
        print("-------------------  ------  ---------")
        for corpus in ALL_CORPORA:
            active = gm.get_active_generation(corpus)
            candidate = gm.get_candidate_generation(corpus)
            cand_str = str(candidate) if candidate else "-"
            print(f"{corpus:20s}  {active:>6d}  {cand_str:>9s}")
    finally:
        store.close()


# ---------------------------------------------------------------------------
# Main / argument parser
# ---------------------------------------------------------------------------


def main(argv: list[str] | None = None) -> None:
    _load_env()

    parser = argparse.ArgumentParser(prog="axiom.rag", description="RAG subsystem CLI")
    parser.add_argument("-v", "--verbose", action="store_true")
    sub = parser.add_subparsers(dest="command")

    # index
    p_index = sub.add_parser("index", help="Index documents into the RAG store")
    p_index.add_argument("paths", nargs="*", help="Paths to index (default: repo docs/)")
    p_index.add_argument(
        "--corpus", default="rag-internal", help="Target corpus (default: rag-internal)"
    )

    # search
    p_search = sub.add_parser("search", help="Search the RAG index")
    p_search.add_argument("query", help="Search query")
    p_search.add_argument("-n", "--limit", type=int, default=5, help="Max results")

    # status
    sub.add_parser("status", help="Show per-corpus index statistics")

    # load-community
    p_lc = sub.add_parser("load-community", help="Load pre-built community corpus")
    p_lc.add_argument(
        "dump_path", nargs="?", default=None, help="Path to .sql dump (default: bundled)"
    )

    # sync
    p_sync = sub.add_parser("sync", help="Sync a corpus from remote source")
    p_sync.add_argument("target", choices=["org"], help="Corpus to sync")

    # watch
    p_watch = sub.add_parser("watch", help="Watch workspace dirs and re-index on change")
    p_watch.add_argument("--corpus", default="rag-internal", help="Target corpus")
    p_watch.add_argument("--quiet", action="store_true", help="Suppress startup output")

    # reindex
    p_reindex = sub.add_parser("reindex", help="Force full re-index of a corpus")
    p_reindex.add_argument("--corpus", default="rag-internal", help="Corpus to reindex")

    # upgrade (build new generation)
    p_upgrade = sub.add_parser("upgrade", help="Build a new RAG generation with semantic chunking")
    p_upgrade.add_argument("--source", required=True, help="Path to source documents")
    p_upgrade.add_argument("--corpus", default="rag-community", help="Target corpus")
    p_upgrade.add_argument("--tier", default="semantic", help="Chunking tier (fixed|semantic)")
    p_upgrade.add_argument("--no-backup", action="store_true", help="Skip pre-upgrade backup")

    # promote / rollback / discard
    p_promote = sub.add_parser("promote", help="Promote a candidate generation to active")
    p_promote.add_argument("generation", type=int, help="Generation number to promote")
    p_promote.add_argument("--corpus", default="rag-community")

    p_rollback = sub.add_parser("rollback", help="Rollback to a previous generation")
    p_rollback.add_argument("generation", type=int, help="Generation to rollback to")
    p_rollback.add_argument("--corpus", default="rag-community")

    p_discard = sub.add_parser("discard", help="Discard a candidate generation")
    p_discard.add_argument("generation", type=int, help="Generation to discard")
    p_discard.add_argument("--corpus", default="rag-community")

    # benchmark
    p_bench = sub.add_parser("benchmark", help="Run A/B quality benchmark")
    p_bench.add_argument(
        "--compare", nargs=2, metavar=("TIER_A", "TIER_B"), default=["fixed", "semantic"]
    )
    p_bench.add_argument("--corpus", default="rag-community")

    # generations
    sub.add_parser("generations", help="Show generation status per corpus")

    # Legacy aliases
    p_ingest = sub.add_parser("ingest", help="[legacy] Alias for index")
    p_ingest.add_argument("--corpus", default="rag-internal")
    p_ingest.add_argument("paths", nargs="*")

    p_stats = sub.add_parser("stats", help="[legacy] Alias for status")
    del p_stats  # suppress unused warning

    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(levelname)s %(name)s: %(message)s",
    )

    dispatch = {
        "index": cmd_index,
        "ingest": cmd_ingest,
        "search": cmd_search,
        "status": cmd_status,
        "stats": cmd_status,
        "load-community": cmd_load_community,
        "sync": cmd_sync,
        "reindex": cmd_reindex,
        "watch": cmd_watch,
        "upgrade": cmd_upgrade,
        "promote": cmd_promote,
        "rollback": cmd_rollback,
        "discard": cmd_discard,
        "benchmark": cmd_benchmark,
        "generations": cmd_generations,
    }

    if args.command in dispatch:
        dispatch[args.command](args)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()
