"""Command-line interface for the neuromind.alzheimers sub-package.

Sub-commands
------------
download   Download PMC articles via pubmed-stream
index      Index downloaded JSON files into ChromaDB
query      Ask a question using the RAG pipeline
compare    Run a question through multiple RAG modes and print a table
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import List


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _setup_logging(verbose: bool) -> None:
    level = logging.DEBUG if verbose else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(levelname)s  %(name)s  %(message)s",
    )


def _print_result_table(results: dict) -> None:
    """Print a compact comparison table for compare mode."""
    from .types import RAGMode

    header = f"{'Mode':<20} {'Papers':>6} {'Ctx tokens':>12} {'Words':>8} {'Time (s)':>10}"
    print("\n" + "=" * len(header))
    print("RAG COMPARISON SUMMARY")
    print("=" * len(header))
    print(header)
    print("-" * len(header))

    mode_papers = {
        RAGMode.NO_RAG: 0,
        RAGMode.RETRIEVAL_ONLY: 5,
        RAGMode.BASIC: 5,
        RAGMode.STANDARD: 10,
        RAGMode.DETAILED: 15,
        RAGMode.COMPREHENSIVE: 25,
    }

    for mode, result in results.items():
        papers = mode_papers.get(mode, "?")
        ctx = result.context_tokens if result.context_tokens else "—"
        words = result.response_words if result.response_words else "—"
        print(
            f"{mode.value:<20} {str(papers):>6} {str(ctx):>12} "
            f"{str(words):>8} {result.total_time:>10.1f}"
        )
    print("=" * len(header))


# ---------------------------------------------------------------------------
# Sub-command handlers
# ---------------------------------------------------------------------------


def cmd_download(args: argparse.Namespace) -> int:
    from .downloader import download_ad_corpus, download_papers

    _setup_logging(args.verbose)

    if args.corpus:
        stats_list = download_ad_corpus(
            max_results_per_topic=args.max_results,
            output_dir=args.output_dir,
            api_key=args.api_key,
            email=args.email,
        )
        total = sum(s.successful for s in stats_list)
        print(f"Corpus download complete — {total} articles across {len(stats_list)} topics.")
        return 0

    if not args.keyword:
        print(
            "ERROR: Provide a KEYWORD or use --corpus to download the full AD corpus.",
            file=sys.stderr,
        )
        return 1

    stats = download_papers(
        keyword=args.keyword,
        max_results=args.max_results,
        output_dir=args.output_dir,
        fmt=args.format,
        api_key=args.api_key,
        email=args.email,
        verbose=args.verbose,
    )
    print(
        f"Downloaded {stats.successful} articles "
        f"({stats.failed} failed, {stats.skipped} skipped) "
        f"in {stats.duration_seconds:.1f}s."
    )
    return 0 if stats.successful > 0 or stats.skipped > 0 else 1


def cmd_index(args: argparse.Namespace) -> int:
    from .indexer import index_directory

    _setup_logging(args.verbose)

    print(f"Indexing {args.papers_dir!r}  →  ChromaDB at {args.db_path!r} …")
    stats = index_directory(
        papers_dir=args.papers_dir,
        db_path=args.db_path,
        collection_name=args.collection,
        chunk_size=args.chunk_size,
        chunk_overlap=args.chunk_overlap,
    )
    print(
        f"Done — {stats.indexed} articles indexed, "
        f"{stats.total_chunks:,} chunks created, "
        f"{stats.skipped} skipped, {stats.failed} failed."
    )
    return 0 if stats.indexed > 0 else 1


def cmd_query(args: argparse.Namespace) -> int:
    from .retriever import Retriever
    from .llm import OllamaLLM
    from .rag import RAGPipeline
    from .types import RAGMode

    _setup_logging(args.verbose)

    try:
        mode = RAGMode(args.mode)
    except ValueError:
        valid = ", ".join(m.value for m in RAGMode)
        print(f"ERROR: Unknown mode {args.mode!r}. Valid modes: {valid}", file=sys.stderr)
        return 1

    retriever = Retriever(db_path=args.db_path, collection_name=args.collection)
    llm = OllamaLLM(model=args.model, host=args.ollama_host, num_ctx=args.num_ctx)
    pipeline = RAGPipeline(retriever=retriever, llm=llm)

    print(f"Query [{mode.value}]: {args.question}\n{'─' * 70}")
    result = pipeline.query(
        question=args.question,
        mode=mode,
        max_response_tokens=args.max_tokens,
        temperature=args.temperature,
    )

    print(result.text)
    print(f"\n{'─' * 70}")
    print(result.summary())

    if args.output:
        out_path = Path(args.output)
        out_path.write_text(
            json.dumps(
                {
                    "query": result.query,
                    "mode": result.mode.value,
                    "text": result.text,
                    "sources": result.sources,
                    "context_tokens": result.context_tokens,
                    "response_words": result.response_words,
                    "retrieve_time": result.retrieve_time,
                    "generate_time": result.generate_time,
                },
                indent=2,
            ),
            encoding="utf-8",
        )
        print(f"Result saved to {args.output}")

    return 0


def cmd_compare(args: argparse.Namespace) -> int:
    from .retriever import Retriever
    from .llm import OllamaLLM
    from .rag import RAGPipeline
    from .types import RAGMode

    _setup_logging(args.verbose)

    modes: List[RAGMode] = []
    for m in args.modes:
        try:
            modes.append(RAGMode(m))
        except ValueError:
            print(f"WARNING: Unknown mode {m!r} — skipping.", file=sys.stderr)

    if not modes:
        modes = list(RAGMode)

    retriever = Retriever(db_path=args.db_path, collection_name=args.collection)
    llm = OllamaLLM(model=args.model, host=args.ollama_host, num_ctx=args.num_ctx)
    pipeline = RAGPipeline(retriever=retriever, llm=llm)

    print(f"Comparing {len(modes)} RAG modes for: {args.question!r}\n")
    results = {}
    for mode in modes:
        print(f"  ⏳ {mode.value} …", end=" ", flush=True)
        result = pipeline.query(
            question=args.question,
            mode=mode,
            max_response_tokens=args.max_tokens,
        )
        results[mode] = result
        print(f"✓ ({result.total_time:.1f}s)")

    _print_result_table(results)

    if args.verbose:
        for mode, result in results.items():
            print(f"\n{'═' * 70}")
            print(f"  {mode.value.upper()}")
            print("═" * 70)
            print(result.text[:1000])
            if len(result.text) > 1000:
                print(f"  … [{len(result.text) - 1000} more chars]")

    return 0


# ---------------------------------------------------------------------------
# Argument parser
# ---------------------------------------------------------------------------


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="alzheimers",
        description="neuromind.alzheimers — Alzheimer's disease RAG research library CLI",
    )
    parser.add_argument("--version", action="version", version="%(prog)s 0.1.1")

    # Shared flags
    shared = argparse.ArgumentParser(add_help=False)
    shared.add_argument("-v", "--verbose", action="store_true", help="Enable DEBUG logging")

    db_shared = argparse.ArgumentParser(add_help=False)
    db_shared.add_argument("--db-path", default="./neuromind_db", metavar="PATH",
                           help="ChromaDB directory (default: ./neuromind_db)")
    db_shared.add_argument("--collection", default="publications", metavar="NAME",
                           help="ChromaDB collection name (default: publications)")

    llm_shared = argparse.ArgumentParser(add_help=False)
    llm_shared.add_argument("--model", default="llama3.2", metavar="TAG",
                            help="Ollama model tag (default: llama3.2)")
    llm_shared.add_argument("--ollama-host", default="http://localhost:11434", metavar="URL",
                            help="Ollama base URL (default: http://localhost:11434)")
    llm_shared.add_argument("--num-ctx", type=int, default=131_072, metavar="N",
                            help="LLM context window size in tokens (default: 131072)")
    llm_shared.add_argument("--max-tokens", type=int, default=3000, metavar="N",
                            help="Maximum generated tokens (default: 3000)")
    llm_shared.add_argument("--temperature", type=float, default=0.3, metavar="T",
                            help="LLM sampling temperature (default: 0.3)")

    sub = parser.add_subparsers(dest="command", required=True)

    # ── download ────────────────────────────────────────────────────────
    p_dl = sub.add_parser("download", parents=[shared],
                          help="Download PubMed Central articles")
    p_dl.add_argument("keyword", nargs="?", default=None,
                      help="Search query (required unless --corpus is set)")
    p_dl.add_argument("--corpus", action="store_true",
                      help="Download the full built-in AD corpus (15 topics)")
    p_dl.add_argument("--max-results", type=int, default=100, metavar="N",
                      help="Max articles per query (default: 100)")
    p_dl.add_argument("-o", "--output-dir", default="./publications", metavar="DIR",
                      help="Output directory (default: ./publications)")
    p_dl.add_argument("--format", choices=["text", "xml", "both"], default="text",
                      help="Article format (default: text)")
    p_dl.add_argument("--api-key", default=None, metavar="KEY",
                      help="NCBI API key (or set NCBI_API_KEY env var)")
    p_dl.add_argument("--email", default=None, metavar="EMAIL",
                      help="Contact email for NCBI (or set NCBI_EMAIL env var)")

    # ── index ────────────────────────────────────────────────────────────
    p_idx = sub.add_parser("index", parents=[shared, db_shared],
                           help="Index downloaded articles into ChromaDB")
    p_idx.add_argument("papers_dir", nargs="?", default="./publications",
                       help="Directory of downloaded JSON files (default: ./publications)")
    p_idx.add_argument("--chunk-size", type=int, default=2000, metavar="N",
                       help="Characters per chunk (default: 2000 ≈ 500 tokens)")
    p_idx.add_argument("--chunk-overlap", type=int, default=200, metavar="N",
                       help="Character overlap between chunks (default: 200)")

    # ── query ────────────────────────────────────────────────────────────
    p_q = sub.add_parser("query", parents=[shared, db_shared, llm_shared],
                         help="Answer a question using the RAG pipeline")
    p_q.add_argument("question", help="Research or clinical question")
    p_q.add_argument("--mode", default="detailed", metavar="MODE",
                     help="RAG mode: no_rag | retrieval_only | basic | standard | "
                          "detailed | comprehensive  (default: detailed)")
    p_q.add_argument("-o", "--output", default=None, metavar="FILE",
                     help="Save result as JSON to this file")

    # ── compare ──────────────────────────────────────────────────────────
    p_cmp = sub.add_parser("compare", parents=[shared, db_shared, llm_shared],
                           help="Compare multiple RAG modes on the same question")
    p_cmp.add_argument("question", help="Research or clinical question")
    p_cmp.add_argument("--modes", nargs="+",
                       default=["no_rag", "basic", "standard", "detailed", "comprehensive"],
                       metavar="MODE",
                       help="Space-separated list of modes to compare")

    return parser


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


def main(argv: list | None = None) -> None:
    parser = _build_parser()
    args = parser.parse_args(argv)

    dispatch = {
        "download": cmd_download,
        "index": cmd_index,
        "query": cmd_query,
        "compare": cmd_compare,
    }
    handler = dispatch.get(args.command)
    if handler is None:
        parser.print_help()
        sys.exit(1)

    sys.exit(handler(args))


if __name__ == "__main__":
    main()
