"""Small command-line entry point; paths are relative to the working directory."""

import json
from pathlib import Path
from typing import Annotated

import typer

from visual_doc_rag import data

app = typer.Typer(no_args_is_help=True, pretty_exceptions_enable=False)
Subset = Annotated[str, typer.Option(help="Dataset subset: hr or industrial.")]
DataDir = Annotated[Path, typer.Option(help="Local data/cache directory; ignored by Git.")]


def check_subset(subset: str) -> None:
    if subset not in data.SUBSETS:
        raise typer.BadParameter(f"Choose one of: {', '.join(data.SUBSETS)}", param_hint="--subset")


@app.command()
def build_demo(
    subset: Subset = "hr",
    data_dir: DataDir = Path("data"),
    output: Annotated[Path | None, typer.Option(help="New bundle directory.")] = None,
    model: Annotated[
        str, typer.Option(help="Model used by the existing sample cache.")
    ] = "gemini-3.5-flash-lite",
    space: Annotated[bool, typer.Option(help="Include Hugging Face Space configuration.")] = False,
) -> None:
    """Package offline retrieval and five cached samples, without calling the provider."""
    from visual_doc_rag.demo import build_demo as pack

    check_subset(subset)
    try:
        result = pack(subset, data_dir, output, model, space)
    except Exception as error:
        typer.echo(f"Demo build failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(f"Demo bundle: {result}")


@app.command()
def download(
    subset: Subset = "hr",
    data_dir: DataDir = Path("data"),
    revision: Annotated[
        str, typer.Option(help="HF branch, tag or commit; resolved to a SHA.")
    ] = "main",
) -> None:
    """Download corpus, queries, qrels and document metadata (no original PDFs)."""
    check_subset(subset)
    try:
        manifest = data.download(subset, data_dir, revision)
    except Exception as error:
        typer.echo(f"Download failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(json.dumps(manifest, indent=2))


@app.command()
def inspect(subset: Subset = "hr", data_dir: DataDir = Path("data")) -> None:
    """Inspect local tables, validate joins and save aggregate statistics/sample pages."""
    check_subset(subset)
    try:
        report = data.inspect(subset, data_dir)
    except Exception as error:
        typer.echo(f"Inspection failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(json.dumps({k: v for k, v in report.items() if k != "features"}, indent=2))
    typer.echo(f"Saved schema and statistics: {data_dir / subset / 'inspection.json'}")
    if not all(report["integrity_checks"].values()):
        typer.echo("Integrity checks failed; review inspection.json before indexing.", err=True)
        raise typer.Exit(1)


@app.command()
def pages(subset: Subset = "hr", data_dir: DataDir = Path("data")) -> None:
    """Export all original page images to a revision-aware PNG cache."""
    from visual_doc_rag.pages import export_pages

    check_subset(subset)
    try:
        result = export_pages(subset, data_dir)
    except Exception as error:
        typer.echo(f"Page export failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(f"Exported {len(result['pages'])} pages; {result['image_bytes'] / 1e6:.1f} MB")


@app.command()
def index(
    subset: Subset = "hr",
    data_dir: DataDir = Path("data"),
    device: Annotated[str, typer.Option(help="auto, cuda or cpu for page encoding.")] = "auto",
    mode: Annotated[str, typer.Option(help="all, text or visual.")] = "all",
    model_revision: Annotated[
        str, typer.Option(help="Adapter revision; reuse cached pin by default.")
    ] = "main",
    base_revision: Annotated[
        str | None, typer.Option(help="Optional pinned base-model revision.")
    ] = None,
) -> None:
    """Build BM25 and visual indexes; reuse completed page embeddings on retry."""
    import torch

    from visual_doc_rag.text_index import TextIndex
    from visual_doc_rag.visual_index import build_visual

    check_subset(subset)
    if mode not in {"all", "text", "visual"} or device not in {"auto", "cuda", "cpu"}:
        raise typer.BadParameter("Use --mode all/text/visual and --device auto/cuda/cpu.")
    torch.set_num_threads(4)
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        source, tables = data.load_tables(subset, data_dir)
        if mode in {"all", "text"}:
            result = TextIndex.build(tables["corpus"], data_dir / subset / "indexes", source)
            typer.echo(f"BM25: {json.dumps(result)}")
        if mode in {"all", "visual"}:
            result = build_visual(subset, data_dir, source, device, model_revision, base_revision)
            typer.echo(json.dumps({k: v for k, v in result.items() if k != "pages"}, indent=2))
    except Exception as error:
        typer.echo(f"Index build failed: {error}", err=True)
        raise typer.Exit(1) from error


@app.command(name="eval")
def eval_command(
    subset: Subset = "hr",
    data_dir: DataDir = Path("data"),
    runs_dir: Annotated[
        Path, typer.Option(help="Directory for timestamped metric reports.")
    ] = Path("runs"),
) -> None:
    """Evaluate text/visual/hybrid on every English query, with CPU latency."""
    import torch

    from visual_doc_rag.evaluate import evaluate, format_table

    check_subset(subset)
    torch.set_num_threads(4)
    try:
        output, report = evaluate(subset, data_dir, runs_dir)
    except Exception as error:
        typer.echo(f"Evaluation failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(format_table(report))
    typer.echo(f"Queries: {report['n_queries']}; pages: {report['n_pages']}")
    typer.echo(f"Saved: {output / 'metrics.json'}")


@app.command()
def search(
    query: Annotated[str, typer.Argument(help="Question or keywords to retrieve pages for.")],
    subset: Subset = "hr",
    data_dir: DataDir = Path("data"),
    method: Annotated[str, typer.Option(help="text, visual or hybrid.")] = "hybrid",
    top_k: Annotated[int, typer.Option(min=1, help="Number of page results to display.")] = 5,
    export: Annotated[
        bool, typer.Option(help="Save an HTML report, page images and JSON under runs/search/.")
    ] = False,
) -> None:
    """Retrieve pages and print local image paths, without generating an answer."""
    from visual_doc_rag.retrieve import PageRetriever

    check_subset(subset)
    if method not in {"text", "visual", "hybrid"} or not query.strip():
        raise typer.BadParameter("Provide a nonempty query and --method text/visual/hybrid.")
    try:
        retriever = PageRetriever(subset, data_dir)
        results = retriever.search(query, method, top_k)
        output = {"method": method, "pages": results}
        if export:
            from visual_doc_rag.search_export import export_search

            output["exported_report"] = str(
                export_search(query, method, retriever.source, results, retriever.tables["corpus"])
            )
    except Exception as error:
        typer.echo(f"Search failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(json.dumps(output, indent=2))


@app.command()
def ask(
    query: Annotated[str, typer.Argument(help="Question to answer using retrieved page images.")],
    subset: Subset = "hr",
    data_dir: DataDir = Path("data"),
    method: Annotated[str, typer.Option(help="text, visual or hybrid.")] = "hybrid",
    export: Annotated[bool, typer.Option(help="Save an answer and evidence HTML report.")] = False,
) -> None:
    """Answer from the top-3 pages with validated citations (requires GEMINI_API_KEY)."""
    from visual_doc_rag.serve import AnswerService

    check_subset(subset)
    if method not in {"text", "visual", "hybrid"} or not query.strip():
        raise typer.BadParameter("Provide a nonempty query and --method text/visual/hybrid.")
    try:
        service = AnswerService(subset, data_dir)
        output = service.ask(query, method)
        if export:
            from visual_doc_rag.search_export import export_search

            output["exported_report"] = str(
                export_search(
                    output["query"],
                    method,
                    service.retriever.source,
                    output["pages"],
                    service.retriever.tables["corpus"],
                    answer=output,
                )
            )
    except Exception as error:
        typer.echo(f"Answer failed: {error}", err=True)
        raise typer.Exit(1) from error
    typer.echo(json.dumps(output, indent=2, ensure_ascii=False))