"""
runner.py - End-to-end pipeline orchestration (full senior-architect QAQC).

Single entry point shared by the CLI and the web API so both run identical
logic. Optionally dumps intermediate artifacts for debugging/tuning.

Stage map (display label -> prompt stage):
    images        -> Stage 0 (pdf_processor)
    extract       -> Stage 2 (extractor, LLM)
    sheet index   -> Stage 1 (sheet_index, LLM)
    jurisdiction  -> Stage 0 (jurisdiction, LLM; runs after we have cover data)
    normalize     -> Stage 3 (normalizer, LLM)
    cluster       -> Stage 4 (clusterer, deterministic)
    conflicts     -> Stage 5 (conflict_checker, LLM)
    qaqc          -> Stage 6 (qaqc_review, LLM)
    code          -> Stage 7 (code_review, LLM + retrieval)
    constructab.  -> Stage 8 (constructability, LLM)
    validate      -> Stage 9 (validator, LLM)
    risk          -> Stage 10 (risk, LLM text-only)
    rfis          -> Stage 11 (rfi, LLM text-only)
    report        -> Stage 12 (report, deterministic)
"""

import os
import json
from typing import Dict, Optional, Callable

from backend.pipeline.pdf_processor import convert_pdf_to_images
from backend.pipeline.extractor import extract_assertions
from backend.text_layer import attach_text_layers, coverage_gaps
from backend.pipeline.sheet_index import classify_sheets, derive_project_meta_from_cover
from backend.pipeline.jurisdiction import run_jurisdiction
from backend.pipeline.normalizer import normalize_assertions, build_project_intelligence
from backend.pipeline.clusterer import cluster_by_location
from backend.pipeline.llm_clusterer import cluster_by_location_llm
from backend import config
from backend.agents.disputes import annotate_clusters
from backend.pipeline.conflict_checker import check_conflicts
from backend.pipeline.qaqc_review import senior_review
from backend.pipeline.code_review import code_review
from backend.pipeline.constructability import constructability_review
from backend.pipeline.validator import dedup_validate
from backend.pipeline.risk import score_and_prioritize
from backend.pipeline.rfi import generate_rfis
from backend.pipeline.report import build_report, to_markdown
from backend.pipeline._stage import validate_issue
from backend.llm import (
    reset_cost, get_cost, set_stage, set_text_backend, set_model_overrides,
)


def run_pipeline(
    pdf_path: str,
    out_dir: Optional[str] = None,
    on_stage: Optional[Callable[[str], None]] = None,
    project_input: Optional[Dict] = None,
    source_name: Optional[str] = None,
    text_local: bool = False,
    vision_model: Optional[str] = None,
    text_model: Optional[str] = None,
) -> Dict:
    """
    Run the full QAQC pipeline on one PDF and return the report dict.

    project_input: optional intake fields (project_name, address, occupancy,
    work_type). Cover-sheet-derived values fill any gaps; intake fields win.

    vision_model / text_model: optional per-run OpenRouter (or local text)
    model overrides from the UI. Blank/None keeps config defaults.

    If out_dir is given, writes conflicts.json, report.md, and the intermediate
    artifacts (assertions.json, clusters.json, and one json per QAQC stage).
    """
    def stage(name: str):
        print(f"\n=== {name} ===")
        set_stage(name)
        if on_stage:
            on_stage(name)

    reset_cost()
    set_text_backend(text_local)
    set_model_overrides(vision_model, text_model)
    if vision_model or text_model:
        print(f"[Runner] model overrides: vision={vision_model or '(default)'} "
              f"text={text_model or '(default)'}")

    try:
        return _run_stages(
            pdf_path, out_dir, stage, project_input, source_name, text_local,
        )
    finally:
        # Don't leak per-run picks into a later overlapping/CLI call.
        set_model_overrides(None, None)
        set_text_backend(False)


def _run_stages(
    pdf_path: str,
    out_dir: Optional[str],
    stage: Callable[[str], None],
    project_input: Optional[Dict],
    source_name: Optional[str],
    text_local: bool,
) -> Dict:
    stage("PDF -> images")
    pages = convert_pdf_to_images(pdf_path)
    text_dir = os.path.join(out_dir, "text") if out_dir else None
    attach_text_layers(pdf_path, pages, text_dir=text_dir)

    stage("Extract assertions")
    sheets = extract_assertions(pages)
    coverage_gaps(pages, sheets)  # classic: log-only recall signal

    stage("Classify sheet index")
    sheet_index = classify_sheets(sheets)

    stage("Jurisdiction profile")
    cover_meta = derive_project_meta_from_cover(sheets, source_name or os.path.basename(pdf_path))
    merged_input = {**cover_meta, **(project_input or {})}
    jurisdiction = run_jurisdiction(merged_input)

    stage("Normalize assertions")
    sheets = normalize_assertions(sheets)

    stage("Build project intelligence (GOIDs + relationships)")
    project_intel = build_project_intelligence(sheets)

    stage(f"Cluster by location ({config.CLUSTERER})")
    if config.CLUSTERER == "llm":
        clusters = cluster_by_location_llm(sheets)
        if not clusters:                      # LLM failed -> fall back, don't lose the run
            print("[Cluster] LLM clustering empty; falling back to deterministic")
            clusters = cluster_by_location(sheets)
    else:
        clusters = cluster_by_location(sheets)

    disputed_count = annotate_clusters(clusters)
    if disputed_count:
        print(f"[Cluster] {disputed_count} cluster(s) carry disputed extracted values")

    stage("Reason over clusters (conflicts)")
    conflicts = check_conflicts(clusters, pages)

    stage("Full-set QAQC review")
    qaqc_issues = senior_review(sheets, clusters, conflicts, sheet_index)

    stage("Code / ADA review")
    code_issues = code_review(jurisdiction, sheets, sheet_index)

    stage("Constructability review")
    construct_issues = constructability_review(sheets, clusters, conflicts)

    stage("Validate & deduplicate")
    conflict_issues = [v for v in (validate_issue(c, "conflict") for c in conflicts) if v]
    all_issues = conflict_issues + qaqc_issues + code_issues + construct_issues
    validated = dedup_validate(all_issues)

    stage("Risk scoring & prioritization")
    prioritized = score_and_prioritize(validated)

    stage("RFI / QAQC comments")
    rfis = generate_rfis(prioritized)

    stage("Build report")
    report = build_report(conflicts, sheets, clusters, source=os.path.basename(pdf_path))
    # Extend the deterministic report with the new QAQC stage outputs.
    report["project_input"] = merged_input
    report["jurisdiction"] = jurisdiction
    report["sheet_index"] = sheet_index
    report["project_intelligence"] = project_intel
    report["validated_issues"] = prioritized
    report["rfis"] = rfis
    report["summary"]["by_stage"] = {
        "conflicts": len(conflicts),
        "qaqc": len(qaqc_issues),
        "code": len(code_issues),
        "constructability": len(construct_issues),
        "validated": len(validated),
        "rfis": len(rfis),
    }
    cost = get_cost()
    report["summary"]["cost_usd"] = round(cost["usd"], 4)
    report["summary"]["llm_calls"] = cost["calls"]
    report["summary"]["cached_calls"] = cost.get("cached", 0)
    report["summary"]["cost_by_stage"] = cost.get("by_stage", {})
    report["summary"]["text_backend"] = "local" if text_local else "openrouter"
    report["summary"]["models_used"] = cost.get("models", {})
    report["summary"]["pipeline_mode"] = "classic"
    print(f"[Runner] LLM cost: ${cost['usd']:.4f} over {cost['calls']} live calls"
          f" ({cost.get('cached', 0)} cached)")

    if out_dir:
        os.makedirs(out_dir, exist_ok=True)
        _dump(out_dir, "assertions.json", sheets)
        _dump(out_dir, "clusters.json", [_cluster_slim(c) for c in clusters])
        _dump(out_dir, "sheet_index.json", sheet_index)
        _dump(out_dir, "jurisdiction.json", jurisdiction)
        _dump(out_dir, "project_intelligence.json", project_intel)
        _dump(out_dir, "qaqc_issues.json", qaqc_issues)
        _dump(out_dir, "code_issues.json", code_issues)
        _dump(out_dir, "constructability.json", construct_issues)
        _dump(out_dir, "validated_issues.json", prioritized)
        _dump(out_dir, "rfis.json", rfis)
        _dump(out_dir, "conflicts.json", report)
        with open(os.path.join(out_dir, "report.md"), "w") as f:
            f.write(to_markdown(report))
        print(f"\n[Runner] Wrote artifacts to {out_dir}")

    return report


def _cluster_slim(c: Dict) -> Dict:
    """Clusters without base64 noise, for artifact dumps."""
    return {k: v for k, v in c.items() if k != "assertions"} | {
        "assertions": [
            {kk: vv for kk, vv in a.items() if kk != "base64"}
            for a in c.get("assertions", [])
        ]
    }


def _dump(out_dir: str, name: str, obj) -> None:
    with open(os.path.join(out_dir, name), "w") as f:
        json.dump(obj, f, indent=2)