orchestrator.py

"""Pipeline orchestrator — chains all 5 agents with timing and error handling."""
from __future__ import annotations
import json
import sys
from collections.abc import Callable
from datetime import datetime, timezone
from pathlib import Path

from agents import ingestion, synthesis, rca, conflict_detector, comms_generator
from models.llm_client import LLMClient
from models.schemas import PipelineResult, SecurityReport, SecurityEventRecord
from infra.clock import now_iso, set_scenario_date
from infra.config import get_config
from infra.logger import get_logger
from infra.security import reset_scanner
from instrumentation.timer import StageTimer
from instrumentation.report import generate_run_report

logger = get_logger("orchestrator")


def run_pipeline(
    scenario_dir: str | Path,
    run_id: str | None = None,
    on_stage: "Callable[[str], None] | None" = None,
    feedback: str | None = None,
    scenario_date: str | None = None,
) -> PipelineResult:
    """
    Run the complete 5-agent escalation synthesis pipeline.

    Args:
        scenario_dir: Path to scenario directory containing artifact JSON files
        run_id: Optional identifier for this run (defaults to scenario+timestamp)

    Returns:
        PipelineResult with all agent outputs and metrics
    """
    config = get_config()
    scenario_path = Path(scenario_dir)
    scenario_id = scenario_path.name

    # Pin the thread-local clock so all agents use the scenario date if provided
    set_scenario_date(scenario_date)

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
    run_id = run_id or f"{scenario_id}_{timestamp}"

    _banner(scenario_id, run_id)

    # ── Fresh security scanner for this run ────────────────────────────
    scanner = reset_scanner()

    timer = StageTimer()

    with LLMClient(config) as llm:
        # ── Agent 1: Ingestion ──────────────────────────────────
        if on_stage: on_stage("ingestion")
        _step_header(1, "Artifact Ingestion", "built-in")
        with timer.measure("ingestion"):
            escalation_context = ingestion.run(scenario_path)
        _step_done(f"Loaded {escalation_context.artifact_count} artifacts",
                   timer.timings["ingestion"])
        _security_summary(scanner)

        # ── Agent 2: Synthesis ──────────────────────────────────
        if on_stage: on_stage("synthesis")
        _step_header(2, "Context Synthesis", config.model_synthesis)
        with timer.measure("synthesis"):
            synthesis_report = synthesis.run(escalation_context, llm, feedback=feedback)
        _step_done(
            f"{len(synthesis_report.confirmed_facts)} confirmed facts  ·  "
            f"{len(synthesis_report.timeline)} timeline events",
            timer.timings["synthesis"],
        )

        # ── Agent 3: Conflict Detection ─────────────────────────
        conflict_report = None
        if config.enable_conflict_detection:
            if on_stage: on_stage("conflict")
            _step_header(3, "Conflict Detection", config.model_conflict)
            _cfd_ok = False
            with timer.measure("conflict_detection"):
                try:
                    conflict_report = conflict_detector.run(
                        escalation_context, synthesis_report, llm
                    )
                    _cfd_ok = True
                except Exception as e:
                    logger.error(f"Conflict detection failed: {e} — continuing")
                    conflict_report = _fallback_conflicts(escalation_context.scenario_id)
            if _cfd_ok:
                _step_done(
                    f"{conflict_report.total_conflicts} conflicts  ·  "
                    f"{conflict_report.unresolved_count} unresolved",
                    timer.timings["conflict_detection"],
                )
            else:
                _step_warn("Conflict detection unavailable — manual review required")
        else:
            _step_header(3, "Conflict Detection", "DISABLED")
            conflict_report = _fallback_conflicts(escalation_context.scenario_id)

        # ── Agent 4: Root Cause Analysis ────────────────────────
        rca_report = None
        if config.enable_rca:
            if on_stage: on_stage("rca")
            _step_header(4, "Root Cause Analysis", config.model_rca)
            _rca_ok = False
            with timer.measure("rca"):
                try:
                    rca_report = rca.run(escalation_context, synthesis_report, conflict_report, llm)
                    _rca_ok = True
                except Exception as e:
                    logger.error(f"RCA agent failed: {e} — continuing without RCA")
                    rca_report = _fallback_rca(escalation_context.scenario_id)
            if _rca_ok:
                _step_done(
                    f"{len(rca_report.causal_chain)}-layer chain  ·  "
                    f"confidence {rca_report.confidence_overall.upper()}",
                    timer.timings["rca"],
                )
                _detail(f"Root cause: {rca_report.root_cause.statement[:90]}…")
            else:
                _step_warn("RCA unavailable — manual review required")
        else:
            _step_header(4, "Root Cause Analysis", "DISABLED")
            rca_report = _fallback_rca(escalation_context.scenario_id)

        # ── Agent 5: Communications ─────────────────────────────
        comms_output = None
        if config.enable_comms_generation:
            if on_stage: on_stage("comms")
            _step_header(5, "Stakeholder Communications", config.model_comms)
            _comms_ok = False
            with timer.measure("comms_generation"):
                try:
                    comms_output = comms_generator.run(
                        escalation_context, synthesis_report, rca_report, conflict_report, llm
                    )
                    _comms_ok = True
                except Exception as e:
                    logger.error(f"Comms generation failed: {e} — continuing")
                    comms_output = _fallback_comms(escalation_context.scenario_id)
            if _comms_ok:
                _step_done(
                    f"3 documents generated  ·  "
                    f"{comms_output.audit_manifest.citation_coverage_pct:.0f}% citation coverage",
                    timer.timings["comms_generation"],
                )
            else:
                _step_warn("Comms generation unavailable — manual review required")
        else:
            _step_header(5, "Comms Generation", "DISABLED")
            comms_output = _fallback_comms(escalation_context.scenario_id)

    # ── Build security report ────────────────────────────────
    _raw_sec = scanner.build_report()
    security_report = SecurityReport(
        total_events       = _raw_sec["total_events"],
        total_redacted     = _raw_sec["total_redacted"],
        injection_attempts = _raw_sec["injection_attempts"],
        redactions         = _raw_sec["redactions"],
        classifications    = _raw_sec["classifications"],
        by_severity        = _raw_sec["by_severity"],
        events             = [
            SecurityEventRecord(
                event_type  = e["type"],
                severity    = e["severity"],
                description = e["description"],
                artifact_id = e["artifact"],
                token       = e["token"],
            )
            for e in _raw_sec["events"]
        ],
    )

    # ── Assemble result ──────────────────────────────────────
    run_metrics = {
        "timing": timer.to_dict(),
        "config": {
            "model_synthesis": config.model_synthesis,
            "model_rca": config.model_rca,
            "model_conflict": config.model_conflict,
            "model_comms": config.model_comms,
            "temperature": config.temperature,
        },
    }

    result = PipelineResult(
        scenario_id=scenario_id,
        run_id=run_id,
        escalation_context=escalation_context,
        synthesis_report=synthesis_report,
        rca_report=rca_report,
        conflict_report=conflict_report,
        communications=comms_output,
        security_report=security_report,
        run_metrics=run_metrics,
    )

    # Print summary
    _print_summary(result, timer)

    # Save outputs
    _save_outputs(result)

    return result


# ── Display helpers ──────────────────────────────────────────────────────────

_W = 66   # output column width (all ASCII — safe on any Windows console)

def _banner(scenario_id: str, run_id: str) -> None:
    from infra.logger import _LOG_FILE
    log_short = "..." + str(_LOG_FILE)[-45:]  # keep last 45 chars so it fits
    sep = "=" * _W
    print(flush=True)
    print(sep, flush=True)
    print(f"  ContextIQ  --  Escalation Context Synthesizer", flush=True)
    print(f"  Scenario : {scenario_id.upper()}", flush=True)
    print(f"  Run ID   : {run_id}", flush=True)
    print(f"  Audit log: {log_short}", flush=True)
    print(sep, flush=True)
    print(flush=True)


def _step_header(n: int, label: str, model: str) -> None:
    print(f"\n  >> Step {n}/5  {label}  [{model}]", flush=True)


def _step_done(detail: str, elapsed: float) -> None:
    from instrumentation.spinner import fmt_time
    print(f"     OK  {detail}  ({fmt_time(elapsed)})", flush=True)


def _step_warn(msg: str) -> None:
    print(f"     !!  {msg}", flush=True)


def _detail(msg: str) -> None:
    print(f"         {msg}", flush=True)


def _security_summary(scanner) -> None:
    """Print a one-line security scan result after ingestion."""
    rep        = scanner.build_report()
    redacted   = rep["total_redacted"]
    injections = rep["injection_attempts"]
    critical   = rep["by_severity"].get("critical", 0)
    high       = rep["by_severity"].get("high", 0)
    classes    = rep["classifications"]

    if injections:
        status = f"BLOCKED  {injections} injection attempt(s)"
    elif critical or high:
        status = f"WARN  {critical + high} high/critical event(s)  |  {redacted} values redacted"
    elif redacted:
        status = f"OK  {redacted} PII value(s) redacted"
    else:
        status = "OK  No PII or secrets detected"

    if classes:
        status += f"  |  classified: {', '.join(classes)}"

    print(f"     [SEC] {status}", flush=True)


def _print_summary(result: PipelineResult, timer: StageTimer) -> None:
    from instrumentation.spinner import fmt_time
    ctx_time  = timer.timings.get("ingestion", 0) + timer.timings.get("synthesis", 0)
    cfd_time  = timer.timings.get("conflict_detection", 0)
    audit_pct = result.communications.audit_manifest.citation_coverage_pct
    kpi_ctx   = "PASS" if ctx_time  < 300 else "FAIL"
    kpi_cfd   = "PASS" if cfd_time  <  30 else "FAIL"
    kpi_audit = "PASS" if audit_pct >= 100 else "FAIL"
    sep       = "=" * _W
    sep2      = "-" * _W

    sec = result.security_report
    sec_detail = f"{sec.total_redacted} PII redacted"
    if sec.injection_attempts:
        sec_detail += f"  |  {sec.injection_attempts} injection attempt(s)"
    if sec.classifications:
        sec_detail += f"  |  classified: {', '.join(sec.classifications)}"

    print(sep, flush=True)
    print("  PIPELINE COMPLETE", flush=True)
    print(sep2, flush=True)
    _row("Run ID",          result.run_id)
    _row("Total time",      fmt_time(timer.total))
    _row("Facts extracted", str(len(result.synthesis_report.confirmed_facts)))
    _row("RCA depth",       f"{len(result.rca_report.causal_chain)} layers  confidence {result.rca_report.confidence_overall.upper()}")
    _row("Conflicts",       f"{result.conflict_report.total_conflicts} found  {result.conflict_report.unresolved_count} unresolved")
    _row("Audit coverage",  f"{audit_pct:.1f}%")
    _row("Security",        sec_detail)
    print(sep2, flush=True)
    _row("KPI  Context assembly",   f"{fmt_time(ctx_time)} / 300s  -->  {kpi_ctx}")
    _row("KPI  Conflict detect",    f"{fmt_time(cfd_time)} / 30s   -->  {kpi_cfd}")
    _row("KPI  Audit trail",        kpi_audit)
    print(sep, flush=True)
    print(flush=True)


def _row(label: str, value: str) -> None:
    print(f"  {label + ':':<28}{value}", flush=True)


def _save_outputs(result: PipelineResult) -> None:
    """Save all pipeline outputs to the outputs directory."""
    output_dir = Path("outputs") / result.run_id
    output_dir.mkdir(parents=True, exist_ok=True)

    # JSON outputs
    (output_dir / "security_report.json").write_text(
        result.security_report.model_dump_json(indent=2), encoding="utf-8"
    )
    (output_dir / "triage_brief.json").write_text(
        result.synthesis_report.model_dump_json(indent=2), encoding="utf-8"
    )
    (output_dir / "rca_chain.json").write_text(
        result.rca_report.model_dump_json(indent=2), encoding="utf-8"
    )
    (output_dir / "conflict_report.json").write_text(
        result.conflict_report.model_dump_json(indent=2), encoding="utf-8"
    )
    (output_dir / "audit_manifest.json").write_text(
        result.communications.audit_manifest.model_dump_json(indent=2), encoding="utf-8"
    )
    (output_dir / "run_metrics.json").write_text(
        json.dumps(result.run_metrics, indent=2), encoding="utf-8"
    )

    # Markdown outputs
    (output_dir / "customer_update.md").write_text(
        result.communications.customer_update, encoding="utf-8"
    )
    (output_dir / "executive_brief.md").write_text(
        result.communications.executive_brief, encoding="utf-8"
    )
    (output_dir / "engineering_handoff.md").write_text(
        result.communications.engineering_handoff, encoding="utf-8"
    )

    from infra.logger import _LOG_FILE
    print(f"  ✓  Outputs saved  →  {output_dir}")
    print(f"  ✓  Audit log      →  {_LOG_FILE}")


def _fallback_rca(scenario_id: str):
    """Return a minimal RCA report when the RCA agent fails."""
    from models.schemas import RCAReport, RCANode
    fallback_node = RCANode(
        layer=0,
        node_id="fallback",
        statement="RCA unavailable — manual review required",
        source_artifact_id="unknown",
        source_entry_id="unknown",
        confidence="low",
    )
    return RCAReport(
        scenario_id=scenario_id,
        causal_chain=[fallback_node],
        dead_branches=[],
        root_cause=fallback_node,
        contributing_factors=["RCA agent unavailable"],
        technical_recommendation="Manual RCA required",
        process_recommendation="Manual RCA required",
        confidence_overall="low",
        rca_timestamp=now_iso(),
    )


def _fallback_conflicts(scenario_id: str):
    """Return an empty conflict report when conflict detection fails."""
    from models.schemas import ConflictReport
    return ConflictReport(
        scenario_id=scenario_id,
        total_conflicts=0,
        conflicts=[],
        unresolved_count=0,
        conflict_timestamp=now_iso(),
    )


def _fallback_comms(scenario_id: str):
    """Return a minimal comms output when comms generation fails."""
    from models.schemas import CommunicationsOutput, AuditManifest
    now = now_iso()
    return CommunicationsOutput(
        scenario_id=scenario_id,
        customer_update="Communications unavailable — manual generation required.",
        executive_brief="Communications unavailable — manual generation required.",
        engineering_handoff="Communications unavailable — manual generation required.",
        audit_manifest=AuditManifest(
            scenario_id=scenario_id,
            total_claims=0,
            cited_claims=0,
            citation_coverage_pct=0.0,
            citations=[],
            manifest_timestamp=now,
        ),
        comms_timestamp=now,
    )


if __name__ == "__main__":
    scenario = sys.argv[1] if len(sys.argv) > 1 else "scenarios/northwind"
    run_pipeline(scenario)