"""
Command-line interface for Quber.

Four commands compose into the extraction pipeline: `document`, `table`, `fuse`,
`figure`. They hand work to each other through files rather than in memory, and
each names its output after the source document's stem, so one stage's
`--output-dir` is the next stage's input directory.

`table`, `fuse` and `figure` each finish by writing a completion marker,
`<base>.tables.complete.json`, `<base>.fuse.complete.json` or
`<base>.figure.complete.json`, as their last upload. The marker carries status
`complete` or `failed`, so a downstream step learns from its existence that the
job's artifacts are whole or that the job failed. `document` writes no marker.
On the cloud path the parse runs on the RunPod worker (`runpod/handler.py`),
which writes the parse's marker itself.

`document` parses the PDF and writes `<base>.docling.json` with
`<base>.confidence.json` and `<base>.cells.json` beside it. All three are always
written, because the later stages reload them together through
`ParseResult.load`. Markdown and split-page HTML are optional extras.

`table` extracts the tables and writes `<base>.tables.json`, plus
`<base>.flags.json` when any table carries a review flag: a corrected cell with
an inspect status, header text left out of the table, or a footnote marker that
could not be tied to a cell (`quber.core.extractors.base.cell_flags`).

`fuse` merges the two into `<base>.unified.json` and `<base>.unified.md`, with
`<base>.cells.json` beside them — the parse's positioned page cells, which the
figure stage's value reconciliation requires next to the parse it is handed. It
also writes the fused tables to `<base>.tables.json`, replacing any already in
the output directory, the fusion report `<base>.fusion.json`, and
`<base>.flags.json` when the tables or the heading review raised a flag. By
default it runs both engines in-process. Given `--artifacts-dir` it reloads
`<base>.{docling,confidence,cells}.json` and `<base>.tables.json` from a prior
run and only fuses, which is how the cloud path splits the work across jobs.

`figure` runs last, after fusion, so its enrichment lands in the document
everything downstream reads. It loads the parse back — `<base>.unified.json`
first, then `<base>.docling.json` — reads the pages holding charts, pictures,
and image-printed tables, and writes the enriched parse back under the filename
it read. Alongside it go `<base>.figures.json` and one `<base>.p<N>.ade.json`
per page scanned, so a re-run reuses a page already paid for instead of
submitting it again. Three more are written only when they have content:
`<base>.scanned-tables.json` when a table was read off a page image,
`<base>.figure-values.json` when `<base>.cells.json` sat beside the parse and
the value step is on — each plotted value reconciled between the scan's reading
and a local read grounded in those cells, with a status and a page box per
value — and `<base>.figure-digest.json` when the scan model is dpt-3, the
default.

`ade`, `analyze`, `batch`, `dual`, and `validate` are standalone and not part of
that chain. `playground` serves the answering web app in `quber.playground`,
which runs the pipeline as subprocesses when a document is uploaded. With no
GPU worker configured it runs `fuse`, then `figure`, then its own ingest. With
one configured it runs `table` locally, sends the parse to the RunPod worker,
then runs `fuse --artifacts-dir`, `figure` and the ingest.
"""

import asyncio
import json
import sys
import tempfile
import time
from importlib.metadata import version
from pathlib import Path
from typing import Optional

import click
from loguru import logger

from quber.agents.factory import ModelConfig, ModelProvider
from quber.agents.llm_client import LLMClient
from quber.files.cache import resolve_document
from quber.files.output import output_sink, write_completion_marker
from quber.settings import get_settings

# `TableInferenceProcessor` and the docling-backed parsers are imported lazily
# inside the commands that use them (`analyze`, `batch`, `document`). Importing
# them at module load would pull the GPU/ML stack into every `quber` invocation,
# including the CPU-only `table` cloud job that needs none of it.


@click.group()
@click.version_option(version=version("quber"))
def cli():
    """Quber - Advanced table extraction and inference system."""
    pass


@cli.command()
@click.argument("pdf_path")
@click.option("--output-dir", "-o", default="output", help="Output directory (default: output)")
@click.option(
    "--provider",
    "-p",
    type=click.Choice([p.value for p in ModelProvider]),
    help="Model provider to use (auto-detected if not specified)",
)
@click.option("--model", "-m", help="Specific model to use (uses provider default if not specified)")
@click.option("--max-concurrent", default=5, type=int, help="Maximum concurrent table analyses (default: 5)")
@click.option(
    "--preset",
    type=click.Choice(["tuned-financial", "legacy"]),
    default="tuned-financial",
    help="DoclingParser configuration preset (default: tuned-financial; 'legacy' = pre-QUE-218)",
)
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def analyze(
    pdf_path: str,
    output_dir: str,
    provider: Optional[str],
    model: Optional[str],
    max_concurrent: int,
    preset: str,
    verbose: bool,
):
    """
    Analyze tables in a PDF document with LLM inference.

    Example:
        quber analyze document.pdf
        quber analyze document.pdf --provider anthropic --model claude-3-opus-20240229
    """
    # Configure logging
    log_level = "DEBUG" if verbose else "INFO"
    logger.remove()
    logger.add(sys.stderr, level=log_level)

    # Load configuration; override the runtime knobs from CLI flags.
    settings = get_settings().model_copy(
        update={"log_level": log_level, "max_concurrent_tables": max_concurrent}
    )

    # Convert provider string to enum if provided
    provider_enum = ModelProvider(provider) if provider else None

    # Create processor
    from quber.processors import TableInferenceProcessor

    try:
        processor = TableInferenceProcessor(
            provider=provider_enum,
            model=model,
            settings=settings,
            preset=preset,
        )
    except ValueError as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

    # Resolve to a local path: local stays local; s3:// materializes through
    # the persistent cache (logs an INFO cache-hit line when already cached).
    doc = resolve_document(pdf_path)

    # Process document
    click.echo(f"\nAnalyzing: {pdf_path}")

    async def run_analysis():
        try:
            results = await processor.process_document(doc)

            # Create output directory
            output_path = Path(output_dir)
            output_path.mkdir(parents=True, exist_ok=True)

            # Generate filenames from PDF basename
            base_name = doc.stem
            json_file = output_path / f"{base_name}_analysis.json"
            summary_file = output_path / f"{base_name}_summary.md"

            # Save files
            processor.save_json(results, str(json_file))
            processor.save_summary(results, str(summary_file))

            return results, json_file, summary_file
        except Exception as e:
            logger.error(f"Analysis failed: {e}")
            raise

    results, json_file, summary_file = asyncio.run(run_analysis())

    # Display summary
    click.echo("\n✓ Analysis complete")
    click.echo(f"  - Found {results['total_tables']} tables")
    click.echo(f"  - JSON output: {json_file}")
    click.echo(f"  - Summary: {summary_file}")


@cli.command()
@click.argument("directory", type=click.Path(exists=True, file_okay=False, dir_okay=True))
@click.option("--output-dir", "-o", default="output", help="Output directory for results (default: output)")
@click.option(
    "--provider", "-p", type=click.Choice([p.value for p in ModelProvider]), help="Model provider to use"
)
@click.option("--model", "-m", help="Specific model to use")
@click.option(
    "--max-concurrent",
    default=5,
    type=int,
    help="Maximum concurrent table analyses per document (default: 5)",
)
@click.option("--pattern", default="*.pdf", help="File pattern to match (default: *.pdf)")
def batch(
    directory: str,
    output_dir: str,
    provider: Optional[str],
    model: Optional[str],
    max_concurrent: int,
    pattern: str,
):
    """
    Process multiple PDF documents in a directory.

    Example:
        quber batch documents/
        quber batch documents/ --pattern "*.pdf" --output-dir results/
    """
    # Load configuration; override max concurrency from the CLI flag.
    settings = get_settings().model_copy(update={"max_concurrent_tables": max_concurrent})

    # Find PDF files
    dir_path = Path(directory)
    pdf_files = list(dir_path.glob(pattern))

    if not pdf_files:
        click.echo(f"No files found matching pattern: {pattern}", err=True)
        sys.exit(1)

    click.echo(f"Found {len(pdf_files)} PDF files to process")

    # Create output directory
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Convert provider string to enum if provided
    provider_enum = ModelProvider(provider) if provider else None

    # Create processor
    from quber.processors import TableInferenceProcessor

    try:
        processor = TableInferenceProcessor(provider=provider_enum, model=model, settings=settings)
    except ValueError as e:
        click.echo(f"Error: {e}", err=True)
        sys.exit(1)

    # Process each file
    async def process_batch():
        results = []
        for pdf_file in pdf_files:
            click.echo(f"\nProcessing: {pdf_file.name}")
            try:
                result = await processor.process_document(str(pdf_file))

                # Save outputs
                base_name = pdf_file.stem
                json_path = output_path / f"{base_name}_analysis.json"
                summary_path = output_path / f"{base_name}_summary.md"

                processor.save_json(result, str(json_path))
                processor.save_summary(result, str(summary_path))

                click.echo(f"  ✓ {result['total_tables']} tables analyzed")
                results.append(result)

            except Exception as e:
                click.echo(f"  ✗ Failed: {e}", err=True)
                logger.error(f"Failed to process {pdf_file}: {e}")

        return results

    results = asyncio.run(process_batch())

    # Summary
    total_documents = len(results)
    total_tables = sum(r["total_tables"] for r in results)

    click.echo("\n✓ Batch processing complete!")
    click.echo(f"  - Processed {total_documents}/{len(pdf_files)} documents")
    click.echo(f"  - Analyzed {total_tables} total tables")
    click.echo(f"  - Results saved to: {output_dir}/")


@cli.command()
@click.option(
    "--provider",
    "-p",
    type=click.Choice([p.value for p in ModelProvider]),
    help="Show models for specific provider only",
)
@click.option("--all", "-a", is_flag=True, help="Show all models (can be very long)")
def models(provider: Optional[str], all: bool):
    """List available models for each provider."""
    from quber.agents import AgentFactory

    factory = AgentFactory(enable_logfire=False)

    # Get available providers (with API keys)
    available = factory.get_available_providers()

    # Filter to specific provider if requested
    if provider:
        providers_to_show = [ModelProvider(provider)]
    else:
        providers_to_show = list(ModelProvider)

    click.echo("Model Providers and Available Models:")
    click.echo("=" * 50)
    click.echo("\nNote: You can use any model name directly, even if not listed.")
    click.echo("The list below shows models recognized by the installed pydantic-ai.\n")

    for p in providers_to_show:
        models = factory.list_available_models(p)
        is_available = p.value in available

        status = "✓ (API key found)" if is_available else "✗ (no API key)"
        click.echo(f"\n{p.value.upper()} {status}")
        click.echo("-" * 40)

        model_list = models[p.value]

        # Show first 10 models by default, all if --all flag is used
        if not all and len(model_list) > 10:
            display_models = model_list[:10]
            remaining = len(model_list) - 10
        else:
            display_models = model_list
            remaining = 0

        for model in display_models:
            default = " (default)" if model == ModelConfig.DEFAULT_MODELS.get(p) else ""
            # Truncate very long model names
            if len(model) > 60:
                model_display = model[:57] + "..."
            else:
                model_display = model
            click.echo(f"  • {model_display}{default}")

        if remaining > 0:
            click.echo(f"  ... and {remaining} more models (use --all to see all)")

    if not available:
        click.echo("\n⚠ No API keys found. Please set environment variables:")
        for p in ModelProvider:
            key_var = ModelConfig.API_KEY_VARS.get(p)
            if key_var:
                click.echo(f"  • {key_var} (for {p.value})")


@cli.group()
def db():
    """Database management commands."""
    pass


@db.command()
@click.option("--drop", is_flag=True, help="Drop existing tables before creating (WARNING: data loss)")
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def init(drop: bool, verbose: bool):
    """Initialize the database schema."""
    from quber.db.connection import get_engine, init_db

    log_level = "DEBUG" if verbose else "INFO"
    logger.remove()
    logger.add(sys.stderr, level=log_level)

    click.echo("Initializing database...")
    try:
        engine = get_engine()
        init_db(engine, drop_all=drop)
        click.echo("✓ Database initialized successfully")
    except Exception as e:
        click.echo(f"✗ Failed to initialize database: {e}", err=True)
        sys.exit(1)


@db.command()
@click.argument("json_path", type=click.Path(exists=True))
@click.option("--generate-embeddings", is_flag=True, help="Generate embeddings during import")
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def import_json(json_path: str, generate_embeddings: bool, verbose: bool):
    """Import a single JSON analysis file into the database."""
    from quber.db.importer import import_json_file

    log_level = "DEBUG" if verbose else "INFO"
    logger.remove()
    logger.add(sys.stderr, level=log_level)

    try:
        doc, table_count = import_json_file(json_path, generate_embeddings=generate_embeddings)
        emb_msg = " with embeddings" if generate_embeddings else ""
        click.echo(f"✓ Imported {doc.filename}: {table_count} tables{emb_msg}")
    except Exception as e:
        click.echo(f"✗ Import failed: {e}", err=True)
        sys.exit(1)


@db.command()
@click.argument("directory", type=click.Path(exists=True, file_okay=False, dir_okay=True))
@click.option("--pattern", default="*_analysis.json", help="File pattern to match (default: *_analysis.json)")
@click.option("--generate-embeddings", is_flag=True, help="Generate embeddings during import")
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def import_dir(directory: str, pattern: str, generate_embeddings: bool, verbose: bool):
    """Import all JSON analysis files from a directory."""
    from quber.db.importer import import_directory

    log_level = "DEBUG" if verbose else "INFO"
    logger.remove()
    logger.add(sys.stderr, level=log_level)

    try:
        docs, total_tables = import_directory(
            directory, pattern=pattern, generate_embeddings=generate_embeddings
        )
        emb_msg = " with embeddings" if generate_embeddings else ""
        click.echo(f"✓ Imported {len(docs)} documents with {total_tables} total tables{emb_msg}")
    except Exception as e:
        click.echo(f"✗ Import failed: {e}", err=True)
        sys.exit(1)


@db.command()
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def stats(verbose: bool):
    """Show database statistics."""
    from quber.db.importer import get_import_stats

    log_level = "DEBUG" if verbose else "INFO"
    logger.remove()
    logger.add(sys.stderr, level=log_level)

    try:
        stats_data = get_import_stats()

        click.echo("\nDatabase Statistics")
        click.echo("=" * 50)
        click.echo(f"Total documents: {stats_data['total_documents']}")
        click.echo(f"Total tables: {stats_data['total_tables']}")
        click.echo(f"Avg tables/doc: {stats_data['avg_tables_per_doc']}")

        if stats_data["documents"]:
            click.echo("\nDocuments:")
            for doc in stats_data["documents"]:
                click.echo(f"  • {doc['filename']}: {doc['tables']} tables, {doc['pages']} pages")

    except Exception as e:
        click.echo(f"✗ Failed to retrieve stats: {e}", err=True)
        sys.exit(1)


@db.command()
@click.option(
    "--provider", type=click.Choice(["local", "openai"]), help="Embedding provider (default: from env)"
)
@click.option("--device", type=click.Choice(["cuda", "cpu"]), help="Device for local model (default: cuda)")
@click.option("--batch-size", default=32, type=int, help="Batch size for processing (default: 32)")
@click.option("--verbose", "-v", is_flag=True, help="Enable verbose logging")
def generate_embeddings(provider: Optional[str], device: Optional[str], batch_size: int, verbose: bool):
    """Generate embeddings for all tables in the database."""
    from quber.db import ExtractedTable, get_embedding_service, get_session

    log_level = "DEBUG" if verbose else "INFO"
    logger.remove()
    logger.add(sys.stderr, level=log_level)

    try:
        click.echo("Initializing embedding service...")
        embedding_service = get_embedding_service(provider=provider, device=device)

        click.echo("Loading tables from database...")
        with get_session() as session:
            tables = session.query(ExtractedTable).all()
            total = len(tables)

            if total == 0:
                click.echo("No tables found in database")
                return

            click.echo(f"Found {total} tables to process")

            # Process in batches
            updated = 0
            for i, table in enumerate(tables, 1):
                if table.llm_title and table.llm_description:
                    # Generate embedding
                    embedding = embedding_service.embed_table_metadata(table.llm_title, table.llm_description)

                    # Update both embedding fields
                    table.title_embedding = embedding.tolist()
                    table.description_embedding = embedding.tolist()
                    updated += 1

                    if i % batch_size == 0:
                        session.commit()
                        click.echo(f"  Processed {i}/{total} tables...")

            # Final commit
            session.commit()
            click.echo(f"✓ Generated embeddings for {updated}/{total} tables")

    except Exception as e:
        click.echo(f"✗ Failed to generate embeddings: {e}", err=True)
        logger.exception(e)
        sys.exit(1)


@cli.command(name="document")
@click.argument("pdf_path")
@click.option("--output-dir", "-o", default="output", help="Output directory (local path or s3:// prefix)")
@click.option(
    "--format",
    "fmt",
    type=click.Choice(["markdown", "json", "both"]),
    default="both",
    help="Output format(s) for the parsed DoclingDocument",
)
@click.option(
    "--html",
    is_flag=True,
    help="Also export html_split_page alongside markdown/json (matches docling CLI --to html_split_page)",
)
@click.option(
    "--image-mode",
    type=click.Choice(["embedded", "referenced", "placeholder"]),
    default="referenced",
    help="How images are written in the markdown export (default: referenced, matching tuned-financial canonical CLI)",
)
@click.option(
    "--preset",
    type=click.Choice(["tuned-financial", "legacy"]),
    default="tuned-financial",
    help="DoclingParser configuration preset (default: tuned-financial; 'legacy' = pre-QUE-218)",
)
@click.option(
    "--device",
    type=click.Choice(["auto", "cuda", "cpu", "mps"]),
    default="auto",
    help="Accelerator device. Default 'auto' resolves to cuda when available, otherwise cpu (so an Intel MacBook 'just works', slower).",
)
def parse(
    pdf_path: str,
    output_dir: str,
    fmt: str,
    html: bool,
    image_mode: str,
    preset: str,
    device: str,
):
    """
    Run the `Parser` path: source PDF -> DoclingDocument.

    Writes the document as markdown and/or JSON. No LLM inference.

    By default uses the canonical `tuned-financial` configuration adopted in
    QUE-218 (TableFormerMode.ACCURATE, DoclingParse backend, OCR enabled,
    picture classification, page_batch_size=32) with `--device auto` resolving
    to CUDA when available (QUE-219). Pass `--preset legacy` to revert to
    the pre-QUE-218 behaviour, or `--device cpu` to force CPU on a GPU host.
    """
    from docling.datamodel.accelerator_options import AcceleratorDevice
    from docling_core.types.doc.base import ImageRefMode

    from quber.core.parsers import parser_for_preset

    parser = parser_for_preset(
        preset,
        accelerator_device=AcceleratorDevice(device),
    )
    # Resolve to a local path before docling: local stays local; s3://
    # materializes through the persistent cache (so docling's str(source)
    # sees the cached local file, not the bare s3:// URI). Logs an INFO
    # cache-hit line when the document is already cached.
    doc = resolve_document(pdf_path)

    click.echo(f"Parsing ({preset}, device={device}): {pdf_path}")
    result = parser.parse(doc)
    document = result.document

    base = doc.stem
    image_ref_mode = ImageRefMode(image_mode)

    # `output_sink` writes locally when `output_dir` is a path, and uploads the
    # artifacts (including any referenced-image sidecar folder) to the bucket
    # when it is an `s3://` prefix.
    with output_sink(output_dir) as out:
        # The document step's artifact contract: docling.json + confidence.json
        # (per-page scores + per-table OCR/native provenance) + cells.json. Always
        # written so the fusion step (and a later cloud job) can reload via
        # ParseResult.load. Markdown and split-page HTML are optional extras.
        artifacts = result.save(out, base)
        click.echo(f"  - JSON: {artifacts['document']}")
        click.echo(f"  - Confidence/provenance: {artifacts['confidence']}")
        click.echo(f"  - Parsed cells: {artifacts['cells']}")

        if fmt in ("markdown", "both"):
            md_path = out / f"{base}.docling.md"
            document.save_as_markdown(md_path, image_mode=image_ref_mode)
            click.echo(f"  - Markdown: {md_path}")
        if html:
            html_path = out / f"{base}.docling.html"
            document.save_as_html(html_path, image_mode=image_ref_mode, split_page_view=True)
            click.echo(f"  - HTML (split-page): {html_path}")

    click.echo(f"  - Pages: {len(document.pages)}")
    click.echo(f"  - Tables: {len(document.tables)}")
    ocr_tables = sum(1 for p in result.table_provenance if p.verdict == "ocr")
    click.echo(f"  - Tables read by OCR (rendered as image): {ocr_tables}")


@cli.command(name="ade")
@click.argument("source")
@click.option("--output-dir", "-o", default="output", help="Output directory (local path or s3:// prefix)")
@click.option(
    "--model",
    default=None,
    help="ADE parse model (default: the generally-available dpt-2; dpt-3 preview names route to the v2 API)",
)
@click.option("--page", type=int, default=None, help="Submit only this 1-based page of a PDF")
def ade(source: str, output_dir: str, model: Optional[str], page: Optional[int]):
    """
    Parse a PDF or image through Landing.AI ADE.

    A standalone extraction path for what docling handles poorly: charts
    (the parse reads a chart's plotted values off the page) and rasterized
    pages with no native text layer. Writes the raw parse response and the
    parse markdown, then a completion marker.

    To read the charts in a document into its parse, use `quber figure`, which
    scans only the pages that hold one.
    """
    from quber.providers.landing import run_parse
    from quber.providers.landing.client import DEFAULT_MODEL

    summary = run_parse(source, output_dir=output_dir, model=model or DEFAULT_MODEL, page=page)
    click.echo(f"ADE parse ({summary.model_version}): {source}" + (f" page {page}" if page else ""))
    for name in summary.artifacts:
        click.echo(f"  - {name}")
    click.echo(f"  - Pages: {summary.pages}")
    click.echo(f"  - Credits consumed: {summary.credits}")


@cli.command(name="figure")
@click.argument("source")
@click.option(
    "--parse",
    required=True,
    help="The parse to enrich: a DoclingDocument JSON file, or a directory holding "
    "one, searched for <base>.unified.json then <base>.docling.json. The enriched "
    "parse is written back under the same filename, so pointing --output-dir at the "
    "directory it came from replaces it",
)
@click.option("--output-dir", "-o", default="output", help="Output directory (local path or s3:// prefix)")
@click.option(
    "--model",
    default=None,
    help="ADE scan model. Defaults to the pinned dpt-3 version; 'dpt-3' selects that "
    "same pin, 'dpt-2' reverts to the previous line, and any exact ADE model name "
    "passes through as given",
)
@click.option(
    "--pages",
    type=click.Choice(["nominated", "all"]),
    default="nominated",
    help="Page breadth: 'nominated' scans the pages the parse points at; 'all' scans "
    "every page, for documents whose content the parse cannot be trusted to point at",
)
def figure(source: str, parse: str, output_dir: str, model: Optional[str], pages: str):
    """
    Read the imagery in a document into its parse.

    Three things a page prints are read by nothing else in the pipeline: the
    values plotted in a chart, the content of any other picture, and a table
    printed as an image, whose cells the parse read off the page itself with
    nothing having checked it since. This reads all three through Landing.AI ADE.

    The parse nominates the pages holding them, and every nominated page is
    scanned; --pages all widens the run to every page of the document. Each page
    is submitted whole, so a chart's title above it and the note at the foot
    of the page come back with it, and a table comes back as a grid with a box on
    every cell, which replaces the parse's reading of it.

    Independent of the rest of the pipeline: in, the parse and the source
    document; out, the same parse enriched, and the page scans. Writes the raw
    response per page scanned, the figure records, the enriched parse under the
    filename it was read from, then a completion marker. A page already scanned
    is read from its stored response instead of being paid for again.

    Run it after `quber fuse` so the enrichment lands in the document the rest
    of the pipeline reads.
    """
    from quber.agents.figure_values import get_figure_value_reader, get_scan_value_parser
    from quber.core.figures.orchestrator import load_document, run_figures_sync, write_marker
    from quber.core.figures.scan import SCAN_DEFAULT_MODEL

    # 'dpt-3' is the CLI spelling of the pinned scan default; the API itself
    # rejects the bare name, so it never passes through by accident.
    if model is None or model.lower() == "dpt-3":
        model = SCAN_DEFAULT_MODEL

    local = resolve_document(source)
    base = local.stem
    document, parse_name = load_document(parse, base)

    # The parse's positioned cells are the fragment source for figure-value
    # reconciliation. They sit beside the parse under <base>.cells.json; a
    # parse directory without them runs the workflow with the values step off.
    parse_dir = Path(parse) if Path(parse).is_dir() else Path(parse).parent
    cells_path = parse_dir / f"{base}.cells.json"
    cells_pages = None
    if cells_path.exists():
        cells_pages = {p["page_no"]: p for p in json.loads(cells_path.read_text(encoding="utf-8"))}
    else:
        logger.warning("No {} beside the parse; figure values skipped", cells_path.name)

    started = time.time()
    try:
        result = run_figures_sync(
            local,
            document,
            base,
            parse_name,
            output_dir=output_dir,
            model=model,
            pages=pages,
            cells_pages=cells_pages,
            value_parser=get_scan_value_parser(),
            value_reader=get_figure_value_reader(),
        )
    except Exception as exc:
        write_marker(
            output_dir,
            base,
            {
                "status": "failed",
                "document": base,
                "stage": "figure",
                "model": model,
                "error": f"{type(exc).__name__}: {exc}",
                "total_seconds": round(time.time() - started, 1),
            },
        )
        raise

    run = result.run
    click.echo(f"Figure scan: {source}")
    for name in result.artifacts:
        click.echo(f"  - {name}")
    click.echo(f"  - Pages nominated: {run.nominated}")
    click.echo(f"  - Pages scanned: {run.scanned} ({run.submitted} submitted this run)")
    click.echo(f"  - Pages returning a figure: {run.with_figures}")
    click.echo(f"  - Figure records: {len(run.figures)}")
    click.echo(f"  - Tables read off a page image: {len(result.tables)}")
    if result.values is not None:
        click.echo(
            f"  - Figure values: {len(result.values.values)} "
            f"({result.values.reconciled} reconciled, {result.values.flagged} flagged)"
        )
    click.echo(f"  - Credits consumed: {run.credits}")
    for err in run.errors:
        click.echo(f"  - ERROR: {err}")

    write_marker(
        output_dir,
        base,
        {
            "status": "complete",
            "document": base,
            "artifacts": result.artifacts,
            # Which model produced the run, off the per-page scan records —
            # two generations share the library, so a completed run says so.
            "model": model,
            "model_version": next((s.version for s in run.scans if s.version), None),
            "pages_nominated": run.nominated,
            "pages_scanned": run.scanned,
            "pages_submitted": run.submitted,
            "pages_returning_a_figure": run.with_figures,
            "figure_records": len(run.figures),
            "scanned_tables": len(result.tables),
            "credit_usage": run.credits,
            "errors": run.errors,
            "total_seconds": round(time.time() - started, 1),
        },
    )


@cli.command(name="table")
@click.argument("pdf_path")
@click.option("--output-dir", "-o", default="output", help="Output directory (local path or s3:// prefix)")
@click.option(
    "--llm-backend",
    type=click.Choice(["cli", "api", "mock"]),
    default=None,
    help="LLM client backend (default: $QUBER_LLM_BACKEND or api)",
)
@click.option(
    "--no-llm",
    is_flag=True,
    help="Skip the LLM structure-correction step (Camelot-only output)",
)
@click.option("--dpi", default=200, type=int, help="Page rasterization DPI")
@click.option(
    "--engine",
    type=click.Choice(["set-of-mark", "camelot-llm", "correspondence"]),
    default="set-of-mark",
    help="Extraction engine (default: set-of-mark). set-of-mark: vision locates each "
    "table and its region, Camelot fills it in-region, grounded correction "
    "cleans structure. camelot-llm and correspondence are DEPRECATED/dormant.",
)
@click.option(
    "--review",
    is_flag=True,
    help="Also write human-reviewable artifacts beside the JSON: a before/after "
    "HTML (source region + bounding box | corrected markdown), an annotated "
    "PDF (Set-of-Mark boxes drawn on the source pages), and — when any table "
    "carries flagged decisions — the flags HTML review queue. Renders "
    "from the tables already extracted, so it adds no extraction or LLM cost.",
)
def extract(
    pdf_path: str,
    output_dir: str,
    llm_backend: Optional[str],
    no_llm: bool,
    dpi: int,
    engine: str,
    review: bool,
):
    """
    Run the `TableExtractor` path: source PDF -> list[ExtractedTable].

    Engines:

    - set-of-mark (default): the page image is the arbiter. The grid locator names
      every table and its region; Camelot extracts each one constrained to its
      box (so it can neither shatter one table into many nor merge many into
      one); the grounded structure-correction LLM cleans headers/spans/merged
      cells without altering a value. One ExtractedTable per visual table.
    - camelot-llm (DEPRECATED): both Camelot flavors in parallel; per-candidate
      is_table classification; LLM unify per page; structure-correction.
    - correspondence (DEPRECATED): detector lists tables, Camelot chunks matched
      by bbox overlap, completeness-audited, with recovery escalation.

    Output is a JSON file of `ExtractedTable` Pydantic models with
    provenance fields.
    """
    from quber.agents.classifier import MockClassifier
    from quber.agents.llm_client import Backend, MockLLMClient, get_llm_client
    from quber.agents.unifier import MockUnifier
    from quber.core.extractors import (
        CamelotCorrespondenceExtractor,
        CamelotLLMTableExtractor,
        TableExtractor,
    )

    backend_lit: Optional[Backend] = llm_backend  # type: ignore[assignment]
    extractor: TableExtractor
    if engine == "set-of-mark":
        if no_llm:
            raise click.UsageError(
                "--no-llm is not supported with --engine set-of-mark; the grounded "
                "correction step is part of the set-of-mark capture."
            )
        from quber.core.extractors.set_of_mark import SetOfMarkExtractor

        extractor = SetOfMarkExtractor(dpi=dpi, llm=get_llm_client(backend_lit))
        logger.info(f"Extracting (SoM): {pdf_path}")
    elif engine == "correspondence":
        if no_llm:
            raise click.UsageError(
                "--no-llm is not supported with --engine correspondence; the "
                "detector/correspondence/completeness flow is LLM-driven end to end."
            )
        extractor = CamelotCorrespondenceExtractor(dpi=dpi, llm=get_llm_client(backend_lit))
        logger.info(f"Extracting (correspondence): {pdf_path}")
    elif no_llm:
        # --no-llm sidelines every LLM stage: classifier default-accepts,
        # unifier pass-through, correction skipped. Output is raw Camelot
        # candidates from both flavors with provenance fields populated.
        extractor = CamelotLLMTableExtractor(
            llm_client=MockLLMClient(),
            classifier=MockClassifier(),
            unifier=MockUnifier(),
            run_llm_correction=False,
            dpi=dpi,
        )
        logger.info(f"Extracting: {pdf_path}")
    else:
        extractor = CamelotLLMTableExtractor(
            llm_client=get_llm_client(backend_lit),
            run_llm_correction=True,
            dpi=dpi,
        )
        logger.info(f"Extracting: {pdf_path}")

    # Resolve to a local path before extraction: local stays local; s3://
    # materializes through the persistent cache. Logs an INFO cache-hit line
    # when the document is already cached.
    source = resolve_document(pdf_path)
    base = source.stem

    # Every run ends with a `<base>.tables.complete.json` marker — the same
    # contract the parse worker follows. Written strictly after the artifacts
    # (its own upload pass), status `complete` or `failed`, so a downstream
    # join can treat marker existence as "these artifacts are whole" and a
    # failure rides the same channel as a success.
    started = time.time()
    stage = "extract"
    try:
        tables = extractor.extract_tables_sync(source)
    except Exception as exc:
        write_completion_marker(
            output_dir,
            f"{base}.tables.complete.json",
            {
                "status": "failed",
                "document": base,
                "stage": stage,
                "error": f"{type(exc).__name__}: {exc}",
                "total_seconds": round(time.time() - started, 1),
            },
        )
        raise

    if engine == "correspondence":
        detected_not_extracted = sum(
            1
            for t in tables
            if t.extraction_record and t.extraction_record.status == "detected_not_extracted"
        )
        incomplete = sum(
            1 for t in tables if t.extraction_record and t.extraction_record.status == "incomplete"
        )
        logger.info(
            f"  - Tables: {len(tables)} "
            f"(detected-not-extracted: {detected_not_extracted}, incomplete: {incomplete})"
        )
    else:
        logger.info(f"  - Tables: {len(tables)}")

    # The run's review queue: every cell whose status is registered for
    # inspection, with full document identity — no reverse attribution
    # from page numbers ever again.
    from quber.core.extractors.base import cell_flags

    flags = cell_flags(tables)

    # `output_sink` writes locally when `output_dir` is a path, and uploads the
    # artifacts to the bucket when it is an `s3://` prefix (cloud job output).
    stage = "write_artifacts"
    try:
        with output_sink(output_dir) as out:
            json_path = out / f"{base}.tables.json"
            json_path.write_text(
                json.dumps([t.model_dump() for t in tables], indent=2, default=str), encoding="utf-8"
            )
            logger.info(f"  - JSON: {json_path}")

            if flags:
                flags_path = out / f"{base}.flags.json"
                flags_path.write_text(json.dumps([f.model_dump() for f in flags], indent=2), encoding="utf-8")
                logger.info(f"  - Flags: {flags_path} ({len(flags)} review item(s))")

            if review:
                # Render straight from the tables already extracted -- no second pass.
                from quber.review import render_annotated_pdf, render_flags_html, render_review_html

                html_path = render_review_html([(source, tables)], out / f"{base}.review.html")
                boxed_path = render_annotated_pdf(tables, source, out / f"{base}.boxed.pdf")
                logger.info(f"  - Review HTML: {html_path}")
                logger.info(f"  - Annotated PDF: {boxed_path}")
                flags_html = render_flags_html([(source, tables)], out / f"{base}.flags.html")
                if flags_html:
                    logger.info(f"  - Flags HTML (review queue): {flags_html}")
    except Exception as exc:
        write_completion_marker(
            output_dir,
            f"{base}.tables.complete.json",
            {
                "status": "failed",
                "document": base,
                "stage": stage,
                "error": f"{type(exc).__name__}: {exc}",
                "total_seconds": round(time.time() - started, 1),
            },
        )
        raise

    write_completion_marker(
        output_dir,
        f"{base}.tables.complete.json",
        {
            "status": "complete",
            "document": base,
            "tables": len(tables),
            "flags": len(flags),
            "total_seconds": round(time.time() - started, 1),
        },
    )


@cli.command()
@click.argument("pdf_path")
@click.option("--output-dir", "-o", default="output", help="Output directory")
@click.option("--dpi", default=200, type=int, help="Page rasterization DPI")
@click.option(
    "--no-classify",
    is_flag=True,
    help="Skip the is_table tagging on Camelot grids (grids still returned, untagged)",
)
def dual(pdf_path: str, output_dir: str, dpi: int, no_classify: bool):
    """
    Run the two first-class flows side by side: source PDF -> DualResult.

    The visual flow looks at each rendered page and lets the grid locator
    own the table count and identity; the Camelot flow pulls cell grids
    straight from the PDF and tags each with an is_table verdict (metadata
    only — it never drops a grid). The two run concurrently and neither is
    matched to the other: this is the un-reconciled input for a later
    cross-reference step. No detector, box repair, recovery, or escalation.

    PDF_PATH may be a local path or an s3:// URI; s3 inputs materialize
    through the persistent cache before extraction.
    """
    from quber.core.extractors.dual import DualFlowExtractor

    source = resolve_document(pdf_path)
    extractor = DualFlowExtractor(dpi=dpi, classify_camelot=not no_classify)
    click.echo(f"Dual extraction: {pdf_path}")
    result = extractor.run_sync(source)

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)
    base = source.stem
    json_path = out / f"{base}.dual.json"
    json_path.write_text(json.dumps(result.model_dump(), indent=2, default=str), encoding="utf-8")

    visual_by_page = result.visual_count_by_page()
    camelot_by_page = result.camelot_count_by_page()
    pages = sorted(set(visual_by_page) | set(camelot_by_page))
    click.echo(f"  - Visual tables: {len(result.visual)}   Camelot grids: {len(result.camelot)}")
    click.echo("  - Per page (visual | camelot):")
    for p in pages:
        click.echo(f"      p{p}: {visual_by_page.get(p, 0)} | {camelot_by_page.get(p, 0)}")
    click.echo(f"  - JSON: {json_path}")


@cli.command()
@click.argument("pdf_path")
@click.option("--output-dir", "-o", default="output", help="Output directory (local path or s3:// prefix)")
@click.option(
    "--artifacts-dir",
    default=None,
    help="Cloud-style standalone fusion: load the document artifacts "
    "(<base>.{docling,confidence,cells}.json) and the table artifact "
    "(<base>.tables.json) from this directory and fuse only, skipping both "
    "engines. Without it, the document and table extractions run in-process.",
)
@click.option(
    "--llm-backend",
    type=click.Choice(["cli", "api", "mock"]),
    default=None,
    help="LLM client backend (default: $QUBER_LLM_BACKEND or api)",
)
@click.option("--dpi", default=200, type=int, help="Page rasterization DPI")
@click.option("--preset", type=click.Choice(["tuned-financial", "legacy"]), default="tuned-financial")
@click.option(
    "--image-mode",
    type=click.Choice(["embedded", "referenced", "placeholder"]),
    default="embedded",
    help="How pictures are written in the unified markdown/HTML. embedded "
    "(default) inlines them base64 so the files are self-contained; referenced "
    "writes a sidecar image folder and links it; placeholder drops them to "
    "'<!-- image -->'.",
)
@click.option(
    "--review",
    is_flag=True,
    help="Also write the human-reviewable control artifacts from both workflows: "
    "the table-side before/after HTML, annotated PDF (Set-of-Mark boxes on the "
    "fused tables), and — when any table carries flagged decisions — "
    "the flags HTML review queue, plus the document-side split-page HTML of the "
    "unified DoclingDocument. Rendered from the result already produced, no "
    "second pass.",
)
def fuse(
    pdf_path: str,
    output_dir: str,
    artifacts_dir: Optional[str],
    llm_backend: Optional[str],
    dpi: int,
    preset: str,
    image_mode: str,
    review: bool,
):
    """
    Fuse the document (docling) and table (Set-of-Mark/Camelot) extractions.

    Two views of the same page enrich each other: the Camelot bodies move into
    the docling spine, regions SoM merged are split, charts and image tables are
    annotated, and tables docling found that SoM missed are surfaced explicitly.

    Two control flows over one fusion core:

    - Local (default): the document and table extractions run in-process, then
      fuse.
    - Standalone (--artifacts-dir): load the two extractions' artifacts from a
      directory (as a cloud job would from S3) and fuse only, no engines.

    Writes the unified DoclingDocument (markdown + JSON), the corrected
    Set-of-Mark/Camelot tables (JSON), and a fusion report (per-region match
    kinds and any surfaced errors). With --review, also writes the table-side
    before/after HTML and annotated PDF and the document-side split-page HTML.
    """
    from quber.agents.llm_client import get_llm_client

    llm = get_llm_client(llm_backend)  # type: ignore[arg-type]
    source = resolve_document(pdf_path)
    base = source.stem

    # Like the extraction jobs, fusion ends with its own completion marker
    # (`<base>.fuse.complete.json`), written strictly after the artifacts.
    # The join step ignores fuse markers, so fusing never re-triggers a join.
    fuse_started = time.time()
    try:
        run_fusion(artifacts_dir, base, source, llm, pdf_path, preset, dpi, image_mode, review, output_dir)
    except Exception as exc:
        write_completion_marker(
            output_dir,
            f"{base}.fuse.complete.json",
            {
                "status": "failed",
                "document": base,
                "stage": "fuse",
                "error": f"{type(exc).__name__}: {exc}",
                "total_seconds": round(time.time() - fuse_started, 1),
            },
        )
        raise
    write_completion_marker(
        output_dir,
        f"{base}.fuse.complete.json",
        {
            "status": "complete",
            "document": base,
            "total_seconds": round(time.time() - fuse_started, 1),
        },
    )


def run_fusion(
    artifacts_dir: Optional[str],
    base: str,
    source: Path,
    llm: LLMClient,
    pdf_path: str,
    preset: str,
    dpi: int,
    image_mode: str,
    review: bool,
    output_dir: str,
):
    """One fusion run, artifacts in to artifacts out. The command wrapper owns the marker."""
    from quber.core.fusion import DocumentFusion, fuse_artifacts

    if artifacts_dir:
        from quber.core.extractors.base import ExtractedTable
        from quber.core.parsers import ParseResult
        from quber.files.cache import resolve_artifacts

        # Local directories pass through; an s3:// prefix materializes the four
        # named artifact files through the persistent cache and reads them there.
        adir = resolve_artifacts(
            artifacts_dir,
            [
                f"{base}.docling.json",
                f"{base}.confidence.json",
                f"{base}.cells.json",
                f"{base}.tables.json",
            ],
        )
        logger.info(f"Fusing from artifacts ({adir}): {pdf_path}")
        parse = ParseResult.load(adir, base)
        som_tables = [ExtractedTable(**t) for t in json.loads((adir / f"{base}.tables.json").read_text())]
        result = asyncio.run(fuse_artifacts(parse, som_tables, source, llm))
    else:
        from quber.core.extractors.set_of_mark import SetOfMarkExtractor

        logger.info(f"Fusing ({preset}): {pdf_path}")
        fusion = DocumentFusion(
            parser=parser_for_preset_lazy(preset),
            extractor=SetOfMarkExtractor(dpi=dpi, llm=llm),
            llm=llm,
        )
        result = fusion.fuse_sync(source)

    from collections import Counter

    from docling_core.types.doc.base import ImageRefMode

    image_ref_mode = ImageRefMode(image_mode)

    kinds = Counter(m.kind for m in result.matches)
    logger.info(f"  - Unified document tables: {len(result.document.tables)}")
    logger.info(f"  - Set-of-Mark/Camelot tables: {len(result.tables)}")
    logger.info(f"  - Match kinds: {dict(kinds)}")
    for err in result.errors:
        logger.error(f"  - {err}")

    with output_sink(output_dir) as out:
        unified_json = out / f"{base}.unified.json"
        unified_json.write_text(json.dumps(result.document.export_to_dict(), indent=2), encoding="utf-8")
        unified_md = out / f"{base}.unified.md"
        result.document.save_as_markdown(unified_md, image_mode=image_ref_mode)

        # The parse's positioned page cells, beside the unified parse. The
        # figure stage reconciles chart values only when this file sits next
        # to the parse it is handed. Without it the figure stage logs a warning
        # and runs with the value step off, and a document ingested through fuse
        # arrives holding no figure values at all.
        if result.parse is not None:
            cells_json = out / f"{base}.cells.json"
            cells_json.write_text(
                json.dumps([p.model_dump() for p in result.parse.pages], indent=2, default=str),
                encoding="utf-8",
            )

        tables_json = out / f"{base}.tables.json"
        tables_json.write_text(
            json.dumps([t.model_dump() for t in result.tables], indent=2, default=str),
            encoding="utf-8",
        )

        report_json = out / f"{base}.fusion.json"
        report_json.write_text(
            json.dumps(
                {
                    "match_kinds": dict(kinds),
                    "matches": [m.model_dump() for m in result.matches],
                    "errors": result.errors,
                },
                indent=2,
                default=str,
            ),
            encoding="utf-8",
        )
        logger.info(f"  - Unified: {unified_json} / {unified_md}")
        logger.info(f"  - Tables: {tables_json}")
        logger.info(f"  - Report: {report_json}")

        from quber.core.extractors.base import cell_flags

        flags = cell_flags(result.tables) + result.heading_flags
        if flags:
            flags_path = out / f"{base}.flags.json"
            flags_path.write_text(json.dumps([f.model_dump() for f in flags], indent=2), encoding="utf-8")
            logger.info(f"  - Flags: {flags_path} ({len(flags)} review item(s))")

        if review:
            # Control artifacts, rendered from the result already produced (no
            # second pass): the table workflow's before/after HTML + annotated
            # PDF over the fused tables, and the document workflow's split-page
            # HTML of the unified document.
            from quber.review import render_annotated_pdf, render_flags_html, render_review_html

            review_html = render_review_html([(source, result.tables)], out / f"{base}.review.html")
            boxed_pdf = render_annotated_pdf(result.tables, source, out / f"{base}.boxed.pdf")
            unified_html = out / f"{base}.unified.html"
            result.document.save_as_html(unified_html, image_mode=image_ref_mode, split_page_view=True)
            logger.info(f"  - Review HTML (tables): {review_html}")
            logger.info(f"  - Annotated PDF: {boxed_pdf}")
            logger.info(f"  - Unified HTML (split-page): {unified_html}")
            flags_html = render_flags_html([(source, result.tables)], out / f"{base}.flags.html")
            if flags_html:
                logger.info(f"  - Flags HTML (review queue): {flags_html}")


def parser_for_preset_lazy(preset: str):
    """Build a parser, importing the GPU/ML stack only when fusion runs."""
    from quber.core.parsers import parser_for_preset

    return parser_for_preset(preset)


@cli.command()
@click.argument("pdf_path", type=click.Path(exists=True))
@click.option(
    "--llm-backend",
    type=click.Choice(["cli", "api", "mock"]),
    default=None,
    help="LLM client backend (default: $QUBER_LLM_BACKEND or api)",
)
@click.option("--dpi", default=200, type=int, help="Page rasterization DPI")
def validate(pdf_path: str, llm_backend: Optional[str], dpi: int):
    """
    Run the Camelot-vs-LLM table-count validation hook.

    Surfaces page-level mismatches between Camelot's detected count and
    the LLM's image-based count. Per plan §7, mismatches are surfaced
    explicitly; we do not silently judge which is correct.
    """
    from pdf2image import convert_from_path

    from quber.agents.classifier import MockClassifier
    from quber.agents.llm_client import Backend, get_llm_client
    from quber.agents.unifier import MockUnifier
    from quber.core.extractors import CamelotLLMTableExtractor, TableExtractor
    from quber.core.validate import camelot_vs_llm_count_sync

    backend_lit: Optional[Backend] = llm_backend  # type: ignore[assignment]
    llm = get_llm_client(backend_lit)
    # Validate compares raw Camelot counts vs LLM image counts; the
    # classifier and unifier would distort the Camelot side, so we
    # sideline them with mocks.
    extractor: TableExtractor = CamelotLLMTableExtractor(
        llm_client=llm,
        classifier=MockClassifier(),
        unifier=MockUnifier(),
        run_llm_correction=False,
        dpi=dpi,
    )

    click.echo(f"Extracting tables: {pdf_path}")
    tables = extractor.extract_tables_sync(Path(pdf_path))

    with tempfile.TemporaryDirectory(prefix="quber-validate-") as tmp:
        tmp_path = Path(tmp)
        images = convert_from_path(pdf_path, dpi=dpi, fmt="png", output_folder=str(tmp_path))
        page_images: list[Path] = []
        for i, img in enumerate(images, start=1):
            p = tmp_path / f"page-{i:04d}.png"
            img.save(p, "PNG")
            page_images.append(p)

        click.echo(f"Running validation across {len(page_images)} pages")
        report = camelot_vs_llm_count_sync(Path(pdf_path), tables, page_images, llm)

    if report.has_mismatches:
        click.echo(f"\nMISMATCHES ({len(report.mismatches)}):")
        for m in report.mismatches:
            click.echo(f"  page {m.page}: camelot={m.camelot_count} llm={m.llm_count} (delta {m.delta:+d})")
    else:
        click.echo("\nNo mismatches detected.")
    if report.errors:
        click.echo(f"\nERRORS ({len(report.errors)}):")
        for e in report.errors:
            click.echo(f"  {e}")


@cli.command()
@click.option("--host", default="127.0.0.1", help="Interface to bind (default: 127.0.0.1)")
@click.option("--port", default=8101, type=int, help="Port to serve on (default: 8101)")
@click.option("--reload", is_flag=True, help="Restart on source changes (development)")
def playground(host: str, port: int, reload: bool):
    """Serve the answering playground web app."""
    # The app module (FastAPI, retrieval, the answer agents) loads in the
    # server process uvicorn starts, not here, so `quber --help` stays light.
    import uvicorn

    uvicorn.run("quber.playground.app:app", host=host, port=port, reload=reload)


# Backwards-compatible aliases. The canonical names say what each command
# produces -- `document` parses a PDF into a DoclingDocument, `table` extracts
# the list of tables. The original `parse`/`extract` names predate that and are
# still referenced in docs and existing scripts, so keep them resolving to the
# same callbacks.
cli.add_command(parse, name="parse")
cli.add_command(extract, name="extract")


def main():
    """Main entry point for the CLI."""
    cli()


if __name__ == "__main__":
    main()