cli.py

#!/usr/bin/env python3
"""
NERVE CLI — Axis product datasheet ingestion and query.

nerve ingest --source ./raw-pdfs
nerve query --model "P3268-SLVE"
nerve stats
"""

import json
import sys
from pathlib import Path

import click

CORPUS_DIR = Path(__file__).parent / "nerve-corpus"


@click.group()
def cli():
    """NERVE — Axis product knowledge extraction pipeline."""
    pass


@cli.command()
@click.option("--source", required=True, type=click.Path(exists=True), help="Source directory of PDF/TXT files")
@click.option("--corpus", default=str(CORPUS_DIR), show_default=True, help="Output corpus directory")
@click.option("--limit", default=0, help="Process at most N products (0 = all)")
@click.option("--no-llm", is_flag=True, default=False, help="Disable LLM fallback even if ANTHROPIC_API_KEY is set")
def ingest(source: str, corpus: str, limit: int, no_llm: bool):
    """Ingest PDFs/TXTs from SOURCE directory into the corpus."""
    import os
    from nerve.ingest import ingest_directory
    from nerve.normalize import normalize_text
    from nerve.extract import extract_product
    from nerve.storage import init_db, save_product, save_index

    if no_llm:
        os.environ.pop("ANTHROPIC_API_KEY", None)

    corpus_path = Path(corpus)
    click.echo(f"Source:  {source}")
    click.echo(f"Corpus:  {corpus_path}")
    click.echo()

    # Ingest raw text
    click.echo("Scanning source directory...")
    products_raw, stats = ingest_directory(source)

    click.echo(f"Found {stats.total} products ({stats.txt_sourced} TXT, "
               f"{stats.pdf_extracted} PDF, {stats.zip_extracted} ZIP)")
    if stats.failed:
        click.echo(f"  Read failures: {len(stats.failed)}", err=True)
    click.echo()

    if limit > 0:
        products_raw = products_raw[:limit]
        click.echo(f"Processing first {limit} products (--limit)\n")

    # Initialize storage
    conn = init_db(corpus_path)

    # Process pipeline
    success = 0
    failed = []
    confidence_sum = 0.0
    low_confidence = []

    with click.progressbar(products_raw, label="Extracting", width=40) as bar:
        for raw in bar:
            try:
                normalized = normalize_text(raw.text)
                product = extract_product(raw.stem, normalized)
                save_product(product, corpus_path, conn)
                success += 1
                conf = product["extraction_confidence"]
                confidence_sum += conf
                if conf < 0.5:
                    low_confidence.append((raw.stem, conf))
            except Exception as e:
                failed.append((raw.stem, str(e)))

    save_index(corpus_path, conn)
    conn.close()

    avg_conf = confidence_sum / success if success else 0.0
    click.echo()
    click.echo("=" * 50)
    click.echo(f"Ingestion complete")
    click.echo(f"  Success:          {success}")
    click.echo(f"  Failed:           {len(failed)}")
    click.echo(f"  Avg confidence:   {avg_conf:.3f}")
    click.echo(f"  Low confidence:   {len(low_confidence)}")
    click.echo(f"  Corpus:           {corpus_path}")

    if failed:
        click.echo("\nFailed products:")
        for stem, err in failed[:10]:
            click.echo(f"  {stem}: {err}", err=True)

    if low_confidence:
        click.echo("\nLow-confidence products (< 0.50):")
        for stem, conf in low_confidence[:10]:
            click.echo(f"  {stem}: {conf:.2f}")


@cli.command()
@click.option("--model", required=True, help="Model key (e.g. P3268-SLVE)")
@click.option("--corpus", default=str(CORPUS_DIR), show_default=True)
@click.option("--field", default=None, help="Print only this top-level field")
def query(model: str, corpus: str, field: str | None):
    """Retrieve specs for a model from the corpus."""
    from nerve.query import CorpusQuery

    try:
        q = CorpusQuery(corpus)
    except FileNotFoundError as e:
        click.echo(str(e), err=True)
        sys.exit(1)

    record = q.get(model)
    q.close()

    if record is None:
        click.echo(f"Model not found: {model}", err=True)
        sys.exit(1)

    if field:
        value = record.get(field)
        click.echo(json.dumps(value, indent=2, ensure_ascii=False))
    else:
        click.echo(json.dumps(record, indent=2, ensure_ascii=False))


@cli.command()
@click.option("--corpus", default=str(CORPUS_DIR), show_default=True)
def stats(corpus: str):
    """Print corpus coverage statistics."""
    import sqlite3

    db_path = Path(corpus) / "index.db"
    if not db_path.exists():
        click.echo("No corpus found. Run 'nerve ingest' first.", err=True)
        sys.exit(1)

    conn = sqlite3.connect(str(db_path))
    from nerve.storage import get_stats

    s = get_stats(conn)
    conn.close()

    click.echo(f"\nNERVE Corpus Statistics")
    click.echo("=" * 40)
    click.echo(f"Total products:   {s['total']}")
    click.echo()

    click.echo("By product type:")
    for ptype, count in sorted(s["by_type"].items(), key=lambda x: -x[1]):
        click.echo(f"  {ptype:<25} {count}")

    click.echo()
    click.echo("Confidence distribution:")
    conf = s["confidence"]
    click.echo(f"  High   (>= 0.85):  {conf['high_ge_0.85']}")
    click.echo(f"  Medium (0.60-0.85): {conf['medium_0.6_to_0.85']}")
    click.echo(f"  Low    (< 0.60):   {conf['low_lt_0.6']}")
    click.echo(f"  Average:            {conf['average']:.3f}")


@cli.command()
@click.option("--chipset", default=None, help="Filter by chipset (e.g. ARTPEC-8)")
@click.option("--type", "product_type", default=None, help="Filter by product type")
@click.option("--series", default=None, help="Filter by series (e.g. P32)")
@click.option("--dlpu", is_flag=True, default=False, help="Only DLPU-capable products")
@click.option("--min-confidence", default=0.0, help="Minimum confidence score")
@click.option("--corpus", default=str(CORPUS_DIR), show_default=True)
def search(chipset, product_type, series, dlpu, min_confidence, corpus):
    """Search/filter products in the corpus."""
    from nerve.query import CorpusQuery

    try:
        q = CorpusQuery(corpus)
    except FileNotFoundError as e:
        click.echo(str(e), err=True)
        sys.exit(1)

    results = q.search(
        chipset=chipset,
        product_type=product_type,
        series=series,
        dlpu=dlpu if dlpu else None,
        min_confidence=min_confidence,
    )
    q.close()

    if not results:
        click.echo("No results found.")
        return

    click.echo(f"Found {len(results)} products:\n")
    header = f"{'Model':<20} {'Series':<8} {'Type':<20} {'Chipset':<12} {'Confidence'}"
    click.echo(header)
    click.echo("-" * len(header))
    for r in results:
        click.echo(
            f"{(r['model'] or ''):<20} "
            f"{(r['series'] or ''):<8} "
            f"{(r['product_type'] or ''):<20} "
            f"{(r['chipset'] or ''):<12} "
            f"{r['extraction_confidence']:.2f}"
        )


if __name__ == "__main__":
    cli()