mosaic/mosaic.py

#!/usr/bin/env python3
"""Mosaic — Intelligence Extraction Platform CLI."""
import sys
import os
import logging
import hashlib
import configparser
from pathlib import Path
from datetime import datetime
from typing import Optional

import click
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.prompt import Prompt
from rich.markdown import Markdown
from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn
from rich import box

# ---------------------------------------------------------------------------
# Paths & constants
# ---------------------------------------------------------------------------
BASE_DIR = Path(__file__).resolve().parent
CONFIG_PATH = BASE_DIR / "config" / "mosaic.conf"
PROFILES_DIR = BASE_DIR / "collectors" / "profiles"
DEFAULT_DB = "mosaic.db"

console = Console()
logger = logging.getLogger("mosaic")

DOMAIN_CHOICES = ["tools", "ttps", "tradecraft", "surveillance", "infrastructure", "entities"]
EXPORT_CATEGORIES = ["tools", "ttps", "tradecraft", "entities", "infrastructure"]
QUERY_MODES = ["default", "emulate", "compare", "brief", "scenario"]

# ---------------------------------------------------------------------------
# Config helpers
# ---------------------------------------------------------------------------

def load_config() -> configparser.ConfigParser:
    """Load mosaic.conf; return a ConfigParser (empty sections on missing file)."""
    cfg = configparser.ConfigParser()
    if CONFIG_PATH.exists():
        cfg.read(str(CONFIG_PATH))
    else:
        logger.debug("Config file not found at %s — using defaults", CONFIG_PATH)
    return cfg


def get_api_key(cfg: configparser.ConfigParser) -> Optional[str]:
    """ANTHROPIC_API_KEY env var takes priority over config file."""
    key = os.environ.get("ANTHROPIC_API_KEY")
    if key:
        return key
    return cfg.get("api", "api_key", fallback=None)


def get_db_path(cfg: configparser.ConfigParser, cli_override: Optional[str] = None) -> str:
    """Resolve database path: CLI flag > config > default."""
    if cli_override:
        return cli_override
    rel = cfg.get("general", "database_path", fallback=DEFAULT_DB)
    return str(BASE_DIR / rel)


def setup_logging(verbose: bool = False):
    """Configure root logger."""
    level = logging.DEBUG if verbose else logging.INFO
    fmt = "%(asctime)s  %(name)-18s  %(levelname)-7s  %(message)s"
    logging.basicConfig(level=level, format=fmt, datefmt="%H:%M:%S")


def get_output_dir(cfg: configparser.ConfigParser) -> str:
    """Resolve the output directory from config."""
    rel = cfg.get("general", "output_dir", fallback="output")
    return str(BASE_DIR / rel)


# ---------------------------------------------------------------------------
# Lazy DB / store helpers (avoid opening DB at import time)
# ---------------------------------------------------------------------------

def open_db(db_path: str):
    """Import and return a MosaicDB instance."""
    from storage.db import MosaicDB
    return MosaicDB(db_path)


def open_jsonl_store(db, output_dir: str):
    """Import and return a JSONLStore instance."""
    from storage.jsonl_store import JSONLStore
    return JSONLStore(db, output_dir=os.path.join(output_dir, "extracted"))


def open_cache(output_dir: str):
    """Import and return a DownloadCache instance."""
    from utils.cache import DownloadCache
    return DownloadCache(base_dir=os.path.join(output_dir, "raw"))


# ---------------------------------------------------------------------------
# Profile helpers
# ---------------------------------------------------------------------------

def load_profiles() -> list[dict]:
    """Scan collectors/profiles/*.yaml and return name+description dicts."""
    profiles = []
    if not PROFILES_DIR.is_dir():
        return profiles
    try:
        import yaml
    except ImportError:
        # Fallback: parse name/description with basic string matching
        return _load_profiles_fallback()

    for yf in sorted(PROFILES_DIR.glob("*.yaml")):
        try:
            data = yaml.safe_load(yf.read_text())
            if isinstance(data, dict):
                profiles.append({
                    "file": yf.name,
                    "name": data.get("name", yf.stem),
                    "description": data.get("description", ""),
                    "base_url": data.get("base_url", ""),
                    "expected_min_count": data.get("expected_min_count", ""),
                })
        except Exception as exc:
            logger.warning("Failed to parse %s: %s", yf.name, exc)
    return profiles


def _load_profiles_fallback() -> list[dict]:
    """Parse YAML profiles without the yaml library (best-effort)."""
    profiles = []
    for yf in sorted(PROFILES_DIR.glob("*.yaml")):
        name = yf.stem
        description = ""
        base_url = ""
        for line in yf.read_text().splitlines():
            line = line.strip()
            if line.startswith("name:"):
                name = line.split(":", 1)[1].strip().strip('"').strip("'")
            elif line.startswith("description:"):
                description = line.split(":", 1)[1].strip().strip('"').strip("'")
            elif line.startswith("base_url:"):
                base_url = line.split(":", 1)[1].strip().strip('"').strip("'")
        profiles.append({
            "file": yf.name,
            "name": name,
            "description": description,
            "base_url": base_url,
            "expected_min_count": "",
        })
    return profiles


# ---------------------------------------------------------------------------
# Click CLI
# ---------------------------------------------------------------------------

@click.group(invoke_without_command=True)
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.option("--url", type=str, default=None, help="Ad-hoc source URL.")
@click.option("--profile", "profile_path", type=click.Path(exists=True), default=None,
              help="Custom YAML profile path.")
@click.option("--dry-run", is_flag=True, default=False, help="Estimate without running.")
@click.option("--retry-failed", is_flag=True, default=False, help="Re-attempt failed docs.")
@click.option("--since", type=str, default=None, help="Only docs after DATE (YYYY-MM-DD).")
@click.option("--date-range", nargs=2, type=str, default=None,
              help="Date window: START END (YYYY-MM-DD).")
@click.option("--since-last-run", is_flag=True, default=False,
              help="Only new/changed since last run.")
@click.option("--update", is_flag=True, default=False, help="Re-check cached URLs.")
@click.option("--db", "db_path", type=click.Path(), default=None, help="Database path.")
@click.option("--verbose", "-v", is_flag=True, default=False, help="Enable debug logging.")
@click.pass_context
def cli(ctx, source, url, profile_path, dry_run, retry_failed, since, date_range,
        since_last_run, update, db_path, verbose):
    """Mosaic -- Intelligence Extraction Platform."""
    setup_logging(verbose)
    cfg = load_config()

    ctx.ensure_object(dict)
    ctx.obj["cfg"] = cfg
    ctx.obj["source"] = source
    ctx.obj["url"] = url
    ctx.obj["profile_path"] = profile_path
    ctx.obj["dry_run"] = dry_run
    ctx.obj["retry_failed"] = retry_failed
    ctx.obj["since"] = since
    ctx.obj["date_range"] = date_range
    ctx.obj["since_last_run"] = since_last_run
    ctx.obj["update"] = update
    ctx.obj["db_path"] = get_db_path(cfg, db_path)
    ctx.obj["verbose"] = verbose
    ctx.obj["output_dir"] = get_output_dir(cfg)

    if ctx.invoked_subcommand is None:
        interactive_menu(ctx)


# ---------------------------------------------------------------------------
# collect
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.option("--url", type=str, default=None, help="Ad-hoc source URL.")
@click.option("--name", type=str, default=None, help="Name for ad-hoc source.")
@click.option("--profile", "profile_path", type=click.Path(exists=True), default=None,
              help="Custom YAML profile path.")
@click.option("--limit", type=int, default=None, help="Maximum documents to fetch.")
@click.option("--since", type=str, default=None, help="Only docs after DATE (YYYY-MM-DD).")
@click.option("--date-range", nargs=2, type=str, default=None,
              help="Date window: START END.")
@click.option("--update", is_flag=True, default=False, help="Re-check cached URLs.")
@click.pass_context
def collect(ctx, source, url, name, profile_path, limit, since, date_range, update):
    """Download documents from a source."""
    source = source or ctx.obj.get("source")
    url = url or ctx.obj.get("url")

    if not source and not url and not profile_path:
        console.print("[yellow]No source specified.[/yellow] Use --source, --url, or --profile.")
        raise SystemExit(1)

    from collectors.profile_loader import load_profile, create_adhoc_profile
    from collectors.custom import CustomSource
    from utils.cache import DownloadCache
    from utils.http_limiter import HTTPRateLimiter

    # Load profile
    try:
        if profile_path:
            profile = load_profile(profile_path)
        elif url:
            pname = name or url.split('/')[2].replace('.', '_')
            profile = create_adhoc_profile(url, pname)
        else:
            profile = load_profile(source)
    except (FileNotFoundError, ValueError) as e:
        console.print(f"[red]Profile error: {e}[/red]")
        raise SystemExit(1)

    label = profile['name']
    console.print(Panel(
        f"[bold cyan]Collect[/bold cyan]\n\n"
        f"  Source : {label} — {profile.get('description', '')}\n"
        f"  Limit  : {limit or 'all'}\n"
        f"  Since  : {since or 'n/a'}\n"
        f"  Range  : {' to '.join(date_range) if date_range else 'n/a'}\n"
        f"  Update : {'yes' if update else 'no'}",
        title="[bold]mosaic collect[/bold]",
        border_style="cyan",
    ))

    db = open_db(ctx.obj["db_path"])
    try:
        cache = DownloadCache(os.path.join(ctx.obj["output_dir"], "raw"))
        cfg = ctx.obj["cfg"]
        http_cfg = {
            'timeout': cfg.get('http', 'timeout', fallback='30'),
            'max_retries': cfg.get('http', 'max_retries', fallback='3'),
            'user_agent': cfg.get('http', 'user_agent', fallback='Mozilla/5.0 (compatible; research-crawler/1.0)'),
        }
        collector = CustomSource(profile, db, cache, config=http_cfg)

        since_date = since
        until_date = None
        if date_range:
            since_date, until_date = date_range

        stats = collector.collect(
            since=since_date, until=until_date,
            limit=limit, update=update,
        )

        table = Table(title="Collection Results", box=box.ROUNDED)
        table.add_column("Metric", style="bold")
        table.add_column("Count", justify="right")
        table.add_row("Discovered", str(stats.get('discovered', 0)))
        table.add_row("Downloaded", str(stats.get('downloaded', 0)))
        table.add_row("Already cached", str(stats.get('cached', 0)))
        table.add_row("Failed", str(stats.get('failed', 0)))
        if update:
            table.add_row("Updated", str(stats.get('updated', 0)))
        console.print(table)
    finally:
        db.close()


# ---------------------------------------------------------------------------
# parse
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.option("--retry-failed", is_flag=True, default=False, help="Re-attempt failed docs.")
@click.pass_context
def parse(ctx, source, retry_failed):
    """Parse downloaded documents into structured text."""
    source = source or ctx.obj.get("source")
    retry_failed = retry_failed or ctx.obj.get("retry_failed", False)

    from parsers.base import ParserManager
    from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn

    console.print(Panel(
        f"[bold green]Parse[/bold green]\n\n"
        f"  Source      : {source or 'all'}\n"
        f"  Retry failed: {'yes' if retry_failed else 'no'}",
        title="[bold]mosaic parse[/bold]",
        border_style="green",
    ))

    db = open_db(ctx.obj["db_path"])
    try:
        cfg = ctx.obj["cfg"]
        parser = ParserManager(
            memory_limit_mb=int(cfg.get('parser', 'subprocess_memory_limit_mb', fallback='512')),
            timeout_seconds=int(cfg.get('parser', 'subprocess_timeout_seconds', fallback='60')),
            min_chars_per_page=int(cfg.get('parser', 'min_chars_per_page', fallback='50')),
            max_pdf_pages=int(cfg.get('parser', 'max_pdf_pages', fallback='200')),
        )

        # Get documents to parse
        if source:
            docs = db.get_documents_by_source(source, status='CACHED')
        else:
            docs = db.get_documents_by_source('', status='CACHED')
            # Get all cached docs across sources
            cursor = db.conn.execute(
                "SELECT * FROM documents WHERE status = 'CACHED' AND parse_status IN ('PENDING', ?)",
                ('PARSE_FAILED' if retry_failed else '__NONE__',)
            )
            from storage.models import Document
            docs = [Document.from_row(row) for row in cursor.fetchall()]

        if not docs:
            console.print("[yellow]No documents to parse.[/yellow]")
            return

        parsed = 0
        failed = 0
        needs_ocr = 0

        with Progress(
            SpinnerColumn(), TextColumn("{task.description}"),
            BarColumn(), TaskProgressColumn(),
        ) as progress:
            task = progress.add_task(f"Parsing {len(docs)} docs", total=len(docs))
            for doc in docs:
                if doc.parse_status in ('PARSED', 'NEEDS_OCR') and not retry_failed:
                    progress.advance(task)
                    continue

                db.update_document_status(doc.id, parse_status='IN_PROGRESS')
                result = parser.parse(doc.raw_path, doc.doc_type, doc.content_hash)

                if result.success:
                    doc.text = result.text
                    doc.title = result.title or doc.title
                    doc.char_count = result.char_count
                    if result.metadata:
                        doc.metadata.update(result.metadata)
                    ps = 'NEEDS_OCR' if result.needs_ocr else 'PARSED'
                    doc.parse_status = ps
                    db.insert_document(doc)
                    db.update_document_status(doc.id, parse_status=ps)
                    if result.needs_ocr:
                        needs_ocr += 1
                    else:
                        parsed += 1
                else:
                    db.update_document_status(doc.id, parse_status='PARSE_FAILED')
                    failed += 1

                progress.advance(task)

        table = Table(title="Parse Results", box=box.ROUNDED)
        table.add_column("Metric", style="bold")
        table.add_column("Count", justify="right")
        table.add_row("Parsed", str(parsed))
        table.add_row("Needs OCR", str(needs_ocr))
        table.add_row("Failed", str(failed))
        console.print(table)
    finally:
        db.close()


# ---------------------------------------------------------------------------
# extract
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.option("--domain", "-d", multiple=True,
              type=click.Choice(DOMAIN_CHOICES, case_sensitive=False),
              help="Extraction domain(s). Repeat for multiple.")
@click.option("--dry-run", is_flag=True, default=False, help="Estimate only.")
@click.option("--since", type=str, default=None, help="Only docs after DATE.")
@click.option("--date-range", nargs=2, type=str, default=None, help="Date window.")
@click.option("--since-last-run", is_flag=True, default=False,
              help="Only new/changed since last run.")
@click.option("--retry-failed", is_flag=True, default=False, help="Re-attempt failed.")
@click.pass_context
def extract(ctx, source, domain, dry_run, since, date_range, since_last_run, retry_failed):
    """Run AI-powered extraction on parsed documents."""
    source = source or ctx.obj.get("source")
    dry_run = dry_run or ctx.obj.get("dry_run", False)
    domains = list(domain) if domain else DOMAIN_CHOICES

    if dry_run:
        _extract_dry_run(ctx.obj["db_path"], source, domains, since, date_range, since_last_run)
        return

    console.print(Panel(
        f"[bold magenta]Extract[/bold magenta]\n\n"
        f"  Source         : {source or 'all'}\n"
        f"  Domains        : {', '.join(domains)}\n"
        f"  Since          : {since or 'n/a'}\n"
        f"  Range          : {' to '.join(date_range) if date_range else 'n/a'}\n"
        f"  Since last run : {'yes' if since_last_run else 'no'}\n"
        f"  Retry failed   : {'yes' if retry_failed else 'no'}",
        title="[bold]mosaic extract[/bold]",
        border_style="magenta",
    ))
    console.print(
        "\n[dim]Extraction not yet implemented — see extractors/[/dim]"
    )


def _extract_dry_run(db_path: str, source: Optional[str], domains: list[str],
                     since: Optional[str], date_range: Optional[tuple],
                     since_last_run: bool):
    """Estimate extraction scope without running anything."""
    from utils.text_utils import estimate_tokens

    db = open_db(db_path)
    try:
        if date_range:
            docs = db.get_documents_by_date_range(source=source,
                                                   since=date_range[0],
                                                   until=date_range[1])
        elif since:
            docs = db.get_documents_by_date_range(source=source, since=since)
        elif since_last_run:
            docs = db.get_unextracted_documents(source=source, since_last_run=True)
        else:
            docs = db.get_unextracted_documents(source=source)

        total_chars = sum(d.char_count for d in docs)
        total_tokens = sum(estimate_tokens(d.text) for d in docs if d.text)

        table = Table(title="Extraction Dry-Run Estimate", box=box.ROUNDED)
        table.add_column("Metric", style="bold")
        table.add_column("Value", justify="right")
        table.add_row("Documents", str(len(docs)))
        table.add_row("Total characters", f"{total_chars:,}")
        table.add_row("Estimated tokens", f"{total_tokens:,}")
        table.add_row("Domains", ", ".join(domains))
        table.add_row("Est. API calls", str(len(docs) * len(domains)))
        console.print(table)
    finally:
        db.close()


# ---------------------------------------------------------------------------
# analyze
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.option("--limit", type=int, default=None, help="Max documents per stage.")
@click.pass_context
def analyze(ctx, source, limit):
    """Full pipeline: collect -> parse -> extract -> report."""
    source = source or ctx.obj.get("source")

    if not source:
        console.print("[yellow]No source specified.[/yellow] Use --source.")
        raise SystemExit(1)

    console.print(Panel(
        f"[bold yellow]Analyze — Full Pipeline[/bold yellow]\n\n"
        f"  Source : {source}\n"
        f"  Limit  : {limit or 'all'}\n\n"
        f"  Stages:\n"
        f"    1. Collect  — download documents\n"
        f"    2. Parse    — extract structured text\n"
        f"    3. Extract  — AI-powered extraction\n"
        f"    4. Report   — generate output",
        title="[bold]mosaic analyze[/bold]",
        border_style="yellow",
    ))
    console.print(
        "\n[dim]Full pipeline not yet implemented — run individual stages.[/dim]"
    )


# ---------------------------------------------------------------------------
# query
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.option("--mode", type=click.Choice(QUERY_MODES, case_sensitive=False),
              default="default", help="Query mode.")
@click.option("--actor", type=str, default=None, help="Filter to actor/group name.")
@click.option("--topic", type=str, default=None, help="Filter to topic.")
@click.pass_context
def query(ctx, source, mode, actor, topic):
    """Interactive Q&A against the intelligence database."""
    source = source or ctx.obj.get("source")
    db_path = ctx.obj["db_path"]

    db = open_db(db_path)
    try:
        _query_loop(db, source=source, mode=mode, actor=actor, topic=topic)
    finally:
        db.close()


def _query_loop(db, source: Optional[str] = None, mode: str = "default",
                actor: Optional[str] = None, topic: Optional[str] = None):
    """Rich-based interactive query loop with FTS5 search."""
    console.print(Panel(
        "[bold cyan]Mosaic Query Console[/bold cyan]\n\n"
        "  Type a search query to find documents via full-text search.\n"
        "  Special commands:\n"
        "    [bold]help[/bold]          — show this message\n"
        "    [bold]exit[/bold] / [bold]quit[/bold]  — leave the console\n"
        "    [bold]mode <name>[/bold]  — switch query mode\n"
        "    [bold]source <name>[/bold] — filter by source",
        title="[bold]mosaic query[/bold]",
        border_style="cyan",
    ))

    if source:
        console.print(f"  [dim]Source filter:[/dim] {source}")
    if mode != "default":
        console.print(f"  [dim]Mode:[/dim] {mode}")
    if actor:
        console.print(f"  [dim]Actor:[/dim] {actor}")
    if topic:
        console.print(f"  [dim]Topic:[/dim] {topic}")
    console.print()

    current_source = source
    current_mode = mode

    while True:
        try:
            user_input = Prompt.ask("[bold cyan]mosaic[/bold cyan]")
        except (EOFError, KeyboardInterrupt):
            console.print("\n[dim]Exiting.[/dim]")
            break

        text = user_input.strip()
        if not text:
            continue

        # --- special commands ---
        lower = text.lower()
        if lower in ("exit", "quit", "q"):
            console.print("[dim]Exiting.[/dim]")
            break

        if lower == "help":
            console.print(
                "  [bold]<query>[/bold]       — full-text search\n"
                "  [bold]mode <name>[/bold]  — switch mode (default/emulate/compare/brief/scenario)\n"
                "  [bold]source <name>[/bold] — set source filter (blank to clear)\n"
                "  [bold]exit / quit[/bold]  — leave"
            )
            continue

        if lower.startswith("mode "):
            new_mode = text[5:].strip().lower()
            if new_mode in QUERY_MODES:
                current_mode = new_mode
                console.print(f"  [dim]Mode set to:[/dim] {current_mode}")
            else:
                console.print(f"  [red]Unknown mode.[/red] Choose from: {', '.join(QUERY_MODES)}")
            continue

        if lower.startswith("source "):
            new_source = text[7:].strip()
            current_source = new_source if new_source else None
            console.print(f"  [dim]Source filter:[/dim] {current_source or '(none)'}")
            continue

        # --- FTS search ---
        results = db.search_documents(text, source=current_source, limit=20)
        if not results:
            console.print("  [dim]No results found.[/dim]")
            continue

        console.print(f"\n  [bold]{len(results)} result(s)[/bold]\n")
        for i, doc in enumerate(results, 1):
            title = doc.title or doc.source_url or doc.id[:16]
            date_str = doc.date or ""
            src = doc.source or ""

            console.print(
                f"  [bold cyan]{i}.[/bold cyan] {title}",
                markup=True, highlight=False,
            )
            console.print(
                f"     [dim]{src}  {date_str}  chars={doc.char_count}[/dim]"
            )
            # Show a text snippet (first 300 chars)
            if doc.text:
                snippet = doc.text[:300].replace("\n", " ")
                console.print(f"     {snippet}", markup=False, highlight=False)
            console.print()


# ---------------------------------------------------------------------------
# export
# ---------------------------------------------------------------------------

@cli.command(name="export")
@click.option("--category", "-c",
              type=click.Choice(EXPORT_CATEGORIES, case_sensitive=False),
              default=None, help="Category to export (default: all).")
@click.option("--format", "fmt", type=str, default="jsonl",
              help="Export format (default: jsonl).")
@click.pass_context
def export_cmd(ctx, category, fmt):
    """Export extracted intelligence to JSONL files."""
    db_path = ctx.obj["db_path"]
    output_dir = ctx.obj["output_dir"]

    if fmt != "jsonl":
        console.print(f"[yellow]Format '{fmt}' not supported yet. Using jsonl.[/yellow]")

    db = open_db(db_path)
    try:
        store = open_jsonl_store(db, output_dir)
        with console.status("[bold green]Exporting..."):
            created = store.export(category=category)

        if created:
            table = Table(title="Exported Files", box=box.ROUNDED)
            table.add_column("File", style="bold")
            table.add_column("Path")
            for fpath in created:
                table.add_row(Path(fpath).name, fpath)
            console.print(table)
        else:
            console.print("[dim]Nothing to export — database is empty for the requested category.[/dim]")
    finally:
        db.close()


# ---------------------------------------------------------------------------
# status
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.pass_context
def status(ctx, source):
    """Show progress dashboard with collection, parse, and extraction stats."""
    source = source or ctx.obj.get("source")
    db_path = ctx.obj["db_path"]

    db = open_db(db_path)
    try:
        _show_status(db, source)
    finally:
        db.close()


def _show_status(db, source: Optional[str] = None):
    """Render status tables."""
    counts = db.get_status_counts(source=source)
    ext_counts = db.get_extraction_counts(source=source)
    tokens = db.get_total_tokens(source=source)

    title_suffix = f" [{source}]" if source else ""

    # --- Collection status ---
    coll_table = Table(title=f"Collection Status{title_suffix}", box=box.ROUNDED)
    coll_table.add_column("Status", style="bold")
    coll_table.add_column("Count", justify="right")
    coll_statuses = ["PENDING", "DOWNLOADING", "CACHED", "DOWNLOAD_FAILED"]
    for s in coll_statuses:
        cnt = counts["collection"].get(s, 0)
        style = "green" if s == "CACHED" else ("red" if "FAIL" in s else "")
        coll_table.add_row(s, f"[{style}]{cnt}[/{style}]" if style else str(cnt))
    coll_table.add_row("[bold]Total[/bold]", f"[bold]{counts['total']}[/bold]")

    # --- Parse status ---
    parse_table = Table(title=f"Parse Status{title_suffix}", box=box.ROUNDED)
    parse_table.add_column("Status", style="bold")
    parse_table.add_column("Count", justify="right")
    parse_statuses = ["PENDING", "IN_PROGRESS", "PARSED", "PARSE_FAILED", "NEEDS_OCR"]
    for s in parse_statuses:
        cnt = counts["parsing"].get(s, 0)
        style = "green" if s == "PARSED" else ("red" if "FAIL" in s else
                                                 ("yellow" if s == "NEEDS_OCR" else ""))
        parse_table.add_row(s, f"[{style}]{cnt}[/{style}]" if style else str(cnt))

    # --- Extraction status per extractor ---
    ext_table = Table(title=f"Extraction Status{title_suffix}", box=box.ROUNDED)
    ext_table.add_column("Extractor", style="bold")
    ext_table.add_column("PENDING", justify="right")
    ext_table.add_column("IN_PROGRESS", justify="right")
    ext_table.add_column("DONE", justify="right", style="green")
    ext_table.add_column("FAILED", justify="right", style="red")
    if ext_counts:
        for ext_name, statuses in sorted(ext_counts.items()):
            ext_table.add_row(
                ext_name,
                str(statuses.get("PENDING", 0)),
                str(statuses.get("IN_PROGRESS", 0)),
                str(statuses.get("DONE", 0)),
                str(statuses.get("FAILED", 0)),
            )
    else:
        ext_table.add_row("[dim]no extractions yet[/dim]", "", "", "", "")

    # --- Token usage ---
    token_table = Table(title=f"Token Usage{title_suffix}", box=box.ROUNDED)
    token_table.add_column("Metric", style="bold")
    token_table.add_column("Value", justify="right")
    token_table.add_row("Input tokens", f"{tokens['input_tokens']:,}")
    token_table.add_row("Output tokens", f"{tokens['output_tokens']:,}")
    token_table.add_row("[bold]Total tokens[/bold]",
                        f"[bold]{tokens['total_tokens']:,}[/bold]")

    console.print()
    console.print(coll_table)
    console.print()
    console.print(parse_table)
    console.print()
    console.print(ext_table)
    console.print()
    console.print(token_table)
    console.print()


# ---------------------------------------------------------------------------
# verify
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--source", type=str, default=None, help="Source profile name.")
@click.pass_context
def verify(ctx, source):
    """Re-hash cached files and report integrity status."""
    source = source or ctx.obj.get("source")
    db_path = ctx.obj["db_path"]
    output_dir = ctx.obj["output_dir"]

    db = open_db(db_path)
    try:
        _run_verify(db, output_dir, source)
    finally:
        db.close()


def _run_verify(db, output_dir: str, source: Optional[str] = None):
    """Verify cached files against stored hashes."""
    cache = open_cache(output_dir)
    doc_hashes = db.get_all_document_hashes()

    if not doc_hashes:
        console.print("[dim]No cached documents to verify.[/dim]")
        return

    verified = 0
    mismatched = 0
    missing = 0
    bad_files: list[str] = []

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TextColumn("{task.completed}/{task.total}"),
        console=console,
    ) as progress:
        task = progress.add_task("Verifying files...", total=len(doc_hashes))

        for doc_id, raw_path, expected_hash in doc_hashes:
            if source:
                # Filter by source if provided — look up in DB
                doc = db.get_document(doc_id)
                if doc and doc.source != source:
                    progress.advance(task)
                    continue

            if not raw_path or not Path(raw_path).exists():
                missing += 1
                progress.advance(task)
                continue

            actual_hash = cache.compute_hash(raw_path)
            if actual_hash == expected_hash:
                verified += 1
            else:
                mismatched += 1
                bad_files.append(raw_path)

            progress.advance(task)

    # Results table
    table = Table(title="Integrity Verification Results", box=box.ROUNDED)
    table.add_column("Status", style="bold")
    table.add_column("Count", justify="right")
    table.add_row("[green]Verified[/green]", str(verified))
    table.add_row("[red]Mismatched[/red]", str(mismatched))
    table.add_row("[yellow]Missing[/yellow]", str(missing))
    table.add_row("[bold]Total checked[/bold]", str(verified + mismatched + missing))
    console.print()
    console.print(table)

    if bad_files:
        console.print("\n[red bold]Mismatched files:[/red bold]")
        for bf in bad_files:
            console.print(f"  [red]- {bf}[/red]", markup=False)
    elif mismatched == 0 and missing == 0:
        console.print("\n[green]All cached files passed integrity checks.[/green]")


# ---------------------------------------------------------------------------
# manual
# ---------------------------------------------------------------------------

@cli.command()
@click.option("--output", "-o", type=click.Path(), default="output/FIELD_MANUAL.md",
              help="Output file path.")
@click.option("--title", type=str, default="Operator's Field Manual",
              help="Manual title.")
@click.pass_context
def manual(ctx, output, title):
    """Generate Markdown field manual from extracted intelligence."""
    from analysis.field_manual import generate_field_manual

    db = open_db(ctx.obj["db_path"])
    try:
        path = generate_field_manual(db, output_path=output, title=title)
        console.print(f"\n[bold green]Field manual generated:[/bold green] {path}")

        # Show stats
        import os
        size = os.path.getsize(path)
        with open(path) as f:
            lines = sum(1 for _ in f)
        console.print(f"  Size: {size:,} bytes ({lines:,} lines)")
    finally:
        db.close()


# ---------------------------------------------------------------------------
# sources
# ---------------------------------------------------------------------------

@cli.command()
@click.pass_context
def sources(ctx):
    """List available source profiles."""
    profiles = load_profiles()

    table = Table(title="Source Profiles", box=box.ROUNDED)
    table.add_column("Name", style="bold cyan")
    table.add_column("Description")
    table.add_column("Base URL", style="dim")
    table.add_column("Expected Docs", justify="right")
    table.add_column("File", style="dim")

    if profiles:
        for p in profiles:
            table.add_row(
                p["name"],
                p["description"],
                p["base_url"],
                str(p["expected_min_count"]) if p["expected_min_count"] else "",
                p["file"],
            )
    else:
        table.add_row("[dim]no profiles found[/dim]", "", "", "", "")

    console.print()
    console.print(table)
    console.print(f"\n  [dim]Profile directory: {PROFILES_DIR}[/dim]")
    console.print("  [dim]Use --profile to load a custom YAML file from any path.[/dim]\n")


# ---------------------------------------------------------------------------
# Interactive menu
# ---------------------------------------------------------------------------

MENU_ITEMS = [
    ("1", "Collect", "Download documents", "collect"),
    ("2", "Parse", "Extract structured text", "parse"),
    ("3", "Extract", "AI-powered extraction", "extract"),
    ("4", "Analyze", "Full pipeline", "analyze"),
    ("5", "Query", "Interactive Q&A", "query"),
    ("6", "Export", "Export to JSONL", "export"),
    ("7", "Status", "Progress dashboard", "status"),
    ("8", "Verify", "Integrity check", "verify"),
    ("9", "Sources", "List profiles", "sources"),
    ("0", "Exit", "", None),
]


def interactive_menu(ctx: click.Context):
    """Display the main interactive menu and dispatch commands."""
    while True:
        lines = []
        for key, label, desc, _ in MENU_ITEMS:
            if key == "0":
                lines.append(f"  [bold red][{key}][/bold red]  {label}")
            else:
                lines.append(f"  [bold cyan][{key}][/bold cyan]  {label:<12} [dim]{desc}[/dim]")

        panel_text = "\n".join(lines)
        console.print()
        console.print(Panel(
            panel_text,
            title="[bold white]MOSAIC[/bold white] [dim]Intelligence Extraction Platform[/dim]",
            border_style="bright_blue",
            padding=(1, 2),
        ))

        try:
            choice = Prompt.ask("\n[bold bright_blue]Select[/bold bright_blue]",
                                choices=[m[0] for m in MENU_ITEMS],
                                show_choices=False)
        except (EOFError, KeyboardInterrupt):
            console.print("\n[dim]Exiting.[/dim]")
            break

        if choice == "0":
            console.print("[dim]Exiting.[/dim]")
            break

        # Find the matching menu item
        for key, label, desc, cmd_name in MENU_ITEMS:
            if key == choice and cmd_name:
                _dispatch_interactive(ctx, cmd_name)
                break


def _dispatch_interactive(ctx: click.Context, cmd_name: str):
    """Invoke a subcommand from the interactive menu."""
    cmd = cli.get_command(ctx, cmd_name)
    if cmd is None:
        console.print(f"[red]Unknown command: {cmd_name}[/red]")
        return

    # Some commands benefit from prompting for a source first
    if cmd_name in ("collect", "parse", "extract", "analyze", "status", "verify"):
        source = _prompt_source_optional()
        if source:
            ctx.obj["source"] = source

    try:
        sub_ctx = click.Context(cmd, parent=ctx, info_name=cmd_name)
        sub_ctx.ensure_object(dict)
        sub_ctx.obj = ctx.obj
        cmd.invoke(sub_ctx)
    except SystemExit:
        pass
    except Exception as exc:
        console.print(f"[red]Error: {exc}[/red]")
        logger.exception("Command '%s' failed", cmd_name)


def _prompt_source_optional() -> Optional[str]:
    """Optionally prompt for a source profile name."""
    profiles = load_profiles()
    if profiles:
        names = [p["name"] for p in profiles]
        console.print(f"\n  [dim]Available sources: {', '.join(names)}[/dim]")

    source = Prompt.ask(
        "  [dim]Source (blank for all)[/dim]",
        default="",
        show_default=False,
    )
    return source.strip() or None


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

def main():
    """Entry point for mosaic CLI."""
    # Ensure we can import project modules regardless of cwd
    if str(BASE_DIR) not in sys.path:
        sys.path.insert(0, str(BASE_DIR))
    cli(standalone_mode=True)


if __name__ == "__main__":
    main()