#!/usr/bin/env python3
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import sys
import time
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


HERE = Path(__file__).resolve().parent
if str(HERE) not in sys.path:
    sys.path.insert(0, str(HERE))

os.environ.setdefault("ITERFACT_OUTPUT_DIR", str(HERE / "generated"))
os.environ.setdefault("ITERFACT_ANALYTICS_DIR", str(HERE / "analytics"))

import server as mcp_server  # noqa: E402


DEFAULT_CASES_PATH = HERE / "regression_cases.json"
DEFAULT_REPORTS_ROOT = HERE / "quality_reports" / "regression"


def load_cases(path: Path) -> list[dict[str, Any]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    cases = payload.get("cases") or []
    if not isinstance(cases, list) or not cases:
        raise ValueError("regression cases file must contain a non-empty 'cases' array.")
    return cases


def filter_cases(cases: list[dict[str, Any]], selected_case_ids: set[str], max_cases: int | None) -> list[dict[str, Any]]:
    filtered = []
    for case in cases:
        case_id = str(case.get("id") or "").strip()
        if selected_case_ids and case_id not in selected_case_ids:
            continue
        filtered.append(case)
    if max_cases:
        filtered = filtered[: max_cases]
    return filtered


def canonicalize_numeric_token(value: Any) -> str:
    text = str(value or "").strip()
    text = text.replace(",", "").replace("$", "").replace("%", "")
    return text


def extract_numeric_tokens(text: str) -> set[str]:
    tokens: set[str] = set()
    for match in re.findall(r"-?\$?\d[\d,]*(?:\.\d+)?%?", text or ""):
        normalized = canonicalize_numeric_token(match)
        if normalized:
            tokens.add(normalized)
    return tokens


def collect_component_engine_ids(plan: dict[str, Any]) -> list[str]:
    seen: list[str] = []
    for section in plan.get("sections", []) or []:
        if not isinstance(section, dict):
            continue
        for component in section.get("components", []) or []:
            if not isinstance(component, dict):
                continue
            engine_id = str(component.get("engine") or "").strip()
            if engine_id and engine_id not in seen:
                seen.append(engine_id)
    return seen


async def wait_for_artifact_completion(artifact_id: str, *, timeout_seconds: int, poll_interval_seconds: float) -> dict[str, Any]:
    deadline = time.time() + timeout_seconds
    latest_state: dict[str, Any] = {}
    while time.time() < deadline:
        latest_state = mcp_server.load_artifact_state(artifact_id)
        artifact_status = str(latest_state.get("artifact_status") or "")
        render_job = latest_state.get("render_job") if isinstance(latest_state.get("render_job"), dict) else {}
        render_status = str(render_job.get("status") or "")
        if artifact_status == "published":
            return latest_state
        if artifact_status == "failed" or render_status == "failed":
            return latest_state
        await asyncio.sleep(poll_interval_seconds)
    raise TimeoutError(f"Timed out waiting for artifact {artifact_id} to complete.")


def evaluate_assertions(case: dict[str, Any], *, response: dict[str, Any], state: dict[str, Any]) -> list[dict[str, Any]]:
    assertions = case.get("assertions") if isinstance(case.get("assertions"), dict) else {}
    plan = state.get("plan") if isinstance(state.get("plan"), dict) else {}
    current = state.get("current") if isinstance(state.get("current"), dict) else {}
    render_validation_summary = current.get("render_validation_summary") if isinstance(current.get("render_validation_summary"), dict) else {}
    visible_text = mcp_server.build_rendered_text_excerpt(render_validation_summary, max_chars=20000)
    visible_text_lower = visible_text.lower()
    visible_numeric_tokens = extract_numeric_tokens(visible_text)
    warning_types = {
        str(warning.get("type") or "")
        for warning in (current.get("warnings") or [])
        if isinstance(warning, dict) and warning.get("type")
    }
    wrapper_engine_id = str(plan.get("wrapper", {}).get("engine") or "")
    component_engine_ids = collect_component_engine_ids(plan)
    narrative_spine = mcp_server.build_narrative_spine_snapshot(plan)
    claim_map = narrative_spine.get("claim_map") or []
    key_entities = narrative_spine.get("key_entities") or []
    section_count = len(plan.get("sections", []) or [])
    quality_score = int(current.get("quality_score") or response.get("quality_score") or 0)

    results: list[dict[str, Any]] = []

    def record(name: str, passed: bool, detail: str) -> None:
        results.append({"name": name, "passed": passed, "detail": detail})

    expected_wrapper = assertions.get("expected_wrapper_engine_id")
    if expected_wrapper is not None:
        passed = wrapper_engine_id == str(expected_wrapper)
        record("expected_wrapper_engine_id", passed, f"expected {expected_wrapper}, got {wrapper_engine_id or 'none'}")

    expected_components = [str(item) for item in (assertions.get("expected_component_engine_ids") or []) if str(item)]
    if expected_components:
        missing = [engine_id for engine_id in expected_components if engine_id not in component_engine_ids]
        record("expected_component_engine_ids", not missing, f"missing: {missing or 'none'}; actual={component_engine_ids}")

    forbidden_components = [str(item) for item in (assertions.get("forbidden_component_engine_ids") or []) if str(item)]
    if forbidden_components:
        present = [engine_id for engine_id in forbidden_components if engine_id in component_engine_ids]
        record("forbidden_component_engine_ids", not present, f"present forbidden: {present or 'none'}")

    min_quality_score = assertions.get("min_quality_score")
    if min_quality_score is not None:
        passed = quality_score >= int(min_quality_score)
        record("min_quality_score", passed, f"expected >= {min_quality_score}, got {quality_score}")

    min_sections = assertions.get("min_section_count")
    if min_sections is not None:
        passed = section_count >= int(min_sections)
        record("min_section_count", passed, f"expected >= {min_sections}, got {section_count}")

    max_sections = assertions.get("max_section_count")
    if max_sections is not None:
        passed = section_count <= int(max_sections)
        record("max_section_count", passed, f"expected <= {max_sections}, got {section_count}")

    absent_warning_types = [str(item) for item in (assertions.get("absent_warning_types") or []) if str(item)]
    if absent_warning_types:
        present = sorted(warning_type for warning_type in absent_warning_types if warning_type in warning_types)
        record("absent_warning_types", not present, f"unexpected warnings: {present or 'none'}")

    required_visible_strings = [str(item) for item in (assertions.get("required_visible_strings") or []) if str(item)]
    if required_visible_strings:
        missing = [item for item in required_visible_strings if item.lower() not in visible_text_lower]
        record("required_visible_strings", not missing, f"missing strings: {missing or 'none'}")

    forbidden_visible_strings = [str(item) for item in (assertions.get("forbidden_visible_strings") or []) if str(item)]
    if forbidden_visible_strings:
        present = [item for item in forbidden_visible_strings if item.lower() in visible_text_lower]
        record("forbidden_visible_strings", not present, f"present forbidden strings: {present or 'none'}")

    required_numeric_tokens = [canonicalize_numeric_token(item) for item in (assertions.get("required_numeric_tokens") or []) if canonicalize_numeric_token(item)]
    if required_numeric_tokens:
        missing = [item for item in required_numeric_tokens if item not in visible_numeric_tokens]
        record("required_numeric_tokens", not missing, f"missing numeric tokens: {missing or 'none'}")

    if assertions.get("require_claim_map"):
        record("require_claim_map", bool(claim_map), f"claim_map entries: {len(claim_map)}")

    if assertions.get("require_key_entities"):
        record("require_key_entities", bool(key_entities), f"key_entities entries: {len(key_entities)}")

    return results


async def run_case(case: dict[str, Any], *, wait_timeout_seconds: int, poll_interval_seconds: float) -> dict[str, Any]:
    started = time.perf_counter()
    started_at = datetime.now(timezone.utc).isoformat()
    telemetry_context = {
        "source_surface": "quality_harness",
        "source_mode": "regression",
        "quality_case_id": case.get("id"),
        "quality_case_title": case.get("title"),
    }

    try:
        response = await mcp_server.generate_document(
            prompt=str(case.get("prompt") or "").strip(),
            template_hint=case.get("template_hint"),
            theme_hint=case.get("theme_hint"),
            data=case.get("data"),
            engine_override=case.get("engine_override"),
            section_count=case.get("section_count"),
            uploaded_content=case.get("uploaded_content"),
            interaction_mode=case.get("interaction_mode"),
            telemetry_context=telemetry_context,
            ctx=None,
        )
        artifact_id = response.get("artifact_id")
        if not artifact_id:
            raise ValueError("generate_document did not return an artifact_id")
        state = mcp_server.load_artifact_state(artifact_id)
        artifact_status = str(state.get("artifact_status") or "")
        if artifact_status != "published":
            state = await wait_for_artifact_completion(
                artifact_id,
                timeout_seconds=wait_timeout_seconds,
                poll_interval_seconds=poll_interval_seconds,
            )

        current = state.get("current") if isinstance(state.get("current"), dict) else {}
        local_path = current.get("local_path")
        local_uri = Path(local_path).resolve().as_uri() if local_path else None
        assertion_results = evaluate_assertions(case, response=response, state=state)
        passed_assertions = sum(1 for result in assertion_results if result.get("passed"))
        total_assertions = len(assertion_results)
        assertion_score = round((passed_assertions / total_assertions) * 100, 1) if total_assertions else 100.0
        duration_ms = int((time.perf_counter() - started) * 1000)
        plan = state.get("plan") if isinstance(state.get("plan"), dict) else {}
        current_warnings = current.get("warnings") or []
        warning_types = sorted(
            {
                str(warning.get("type"))
                for warning in current_warnings
                if isinstance(warning, dict) and warning.get("type")
            }
        )
        return {
            "status": "passed" if all(result.get("passed") for result in assertion_results) else "failed_assertions",
            "started_at": started_at,
            "duration_ms": duration_ms,
            "case_id": case.get("id"),
            "case_title": case.get("title"),
            "artifact_id": artifact_id,
            "title": response.get("title"),
            "summary": response.get("summary"),
            "preview_url": current.get("preview_url") or response.get("preview_url"),
            "artifact_page_url": current.get("preview_url") or response.get("preview_url") or local_uri,
            "local_path": local_path,
            "local_uri": local_uri,
            "quality_score": int(current.get("quality_score") or response.get("quality_score") or 0),
            "warning_types": warning_types,
            "wrapper_engine_id": str(plan.get("wrapper", {}).get("engine") or ""),
            "component_engine_ids": collect_component_engine_ids(plan),
            "section_count": len(plan.get("sections", []) or []),
            "section_titles": [section.get("title") for section in (plan.get("sections") or []) if isinstance(section, dict) and section.get("title")],
            "claim_map_size": len(mcp_server.build_narrative_spine_snapshot(plan).get("claim_map") or []),
            "key_entity_count": len(mcp_server.build_narrative_spine_snapshot(plan).get("key_entities") or []),
            "assertion_score": assertion_score,
            "assertions": assertion_results,
            "notes": response.get("notes") or [],
        }
    except Exception as exc:  # noqa: BLE001
        duration_ms = int((time.perf_counter() - started) * 1000)
        return {
            "status": "generation_failed" if "Render validation failed" in str(exc) else "error",
            "started_at": started_at,
            "duration_ms": duration_ms,
            "case_id": case.get("id"),
            "case_title": case.get("title"),
            "error_type": type(exc).__name__,
            "error": str(exc),
        }


def build_summary(results: list[dict[str, Any]]) -> dict[str, Any]:
    passed = [result for result in results if result.get("status") == "passed"]
    failed_assertions = [result for result in results if result.get("status") == "failed_assertions"]
    generation_failed = [result for result in results if result.get("status") == "generation_failed"]
    errored = [result for result in results if result.get("status") == "error"]
    warning_counter: Counter[str] = Counter()

    for result in results:
        warning_counter.update(result.get("warning_types") or [])

    return {
        "total_cases": len(results),
        "passed": len(passed),
        "failed_assertions": len(failed_assertions),
        "generation_failed": len(generation_failed),
        "errors": len(errored),
        "avg_quality_score": round(
            sum(int(result.get("quality_score") or 0) for result in results if result.get("quality_score") is not None)
            / max(1, len([result for result in results if result.get("quality_score") is not None])),
            1,
        )
        if any(result.get("quality_score") is not None for result in results)
        else None,
        "avg_assertion_score": round(
            sum(float(result.get("assertion_score") or 0) for result in results if result.get("assertion_score") is not None)
            / max(1, len([result for result in results if result.get("assertion_score") is not None])),
            1,
        )
        if any(result.get("assertion_score") is not None for result in results)
        else None,
        "warning_types": dict(sorted(warning_counter.items())),
    }


def build_markdown_summary(report: dict[str, Any]) -> str:
    summary = report["summary"]
    lines = [
        "# IterFact MCP Regression Harness",
        "",
        f"- Run at: `{report['generated_at']}`",
        f"- Total cases: `{summary['total_cases']}`",
        f"- Passed: `{summary['passed']}`",
        f"- Failed assertions: `{summary['failed_assertions']}`",
        f"- Generation failed: `{summary['generation_failed']}`",
        f"- Errors: `{summary['errors']}`",
        f"- Avg quality score: `{summary['avg_quality_score']}`",
        f"- Avg assertion score: `{summary['avg_assertion_score']}`",
        "",
        "## Warning types",
    ]
    if summary["warning_types"]:
        for warning_type, count in summary["warning_types"].items():
            lines.append(f"- `{warning_type}`: `{count}`")
    else:
        lines.append("- none")

    lines.append("")
    lines.append("## Case results")
    for result in report["results"]:
        lines.append(
            f"- `{result['case_id']}` {result['case_title']}: "
            f"`{result['status']}`, quality `{result.get('quality_score')}`, assertion score `{result.get('assertion_score')}`"
        )
        for assertion in result.get("assertions") or []:
            marker = "PASS" if assertion.get("passed") else "FAIL"
            lines.append(f"  - `{marker}` {assertion.get('name')}: {assertion.get('detail')}")
        if result.get("error"):
            lines.append(f"  - `ERROR` {result.get('error')}")

    lines.append("")
    lines.append("## Report files")
    lines.append(f"- JSON: `{report['report_path']}`")
    lines.append(f"- Markdown: `{report['summary_path']}`")
    return "\n".join(lines) + "\n"


async def main() -> int:
    parser = argparse.ArgumentParser(description="Run a fixed IterFact regression corpus with explicit assertions.")
    parser.add_argument("--cases", default=str(DEFAULT_CASES_PATH), help="Path to the regression cases JSON file.")
    parser.add_argument("--case", action="append", dest="case_ids", default=[], help="Case id to include. Repeatable.")
    parser.add_argument("--max-cases", type=int, default=None, help="Limit the number of cases.")
    parser.add_argument("--wait-timeout-seconds", type=int, default=180, help="How long to wait for async artifact completion.")
    parser.add_argument("--poll-interval-seconds", type=float, default=1.5, help="How often to poll local artifact state.")
    parser.add_argument("--output-root", default=str(DEFAULT_REPORTS_ROOT), help="Directory where regression reports should be written.")
    args = parser.parse_args()

    cases_path = Path(args.cases).expanduser().resolve()
    output_root = Path(args.output_root).expanduser().resolve()
    output_root.mkdir(parents=True, exist_ok=True)

    cases = filter_cases(load_cases(cases_path), set(args.case_ids), args.max_cases)
    if not cases:
        raise ValueError("No regression cases matched the requested filters.")

    run_stamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    report_dir = output_root / run_stamp
    report_dir.mkdir(parents=True, exist_ok=True)

    results: list[dict[str, Any]] = []
    for case in cases:
        results.append(
            await run_case(
                case,
                wait_timeout_seconds=args.wait_timeout_seconds,
                poll_interval_seconds=args.poll_interval_seconds,
            )
        )

    report_path = report_dir / "regression_report.json"
    summary_path = report_dir / "regression_summary.md"
    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "cases_path": str(cases_path),
        "report_dir": str(report_dir),
        "report_path": str(report_path),
        "summary_path": str(summary_path),
        "case_count": len(cases),
        "cases": cases,
        "results": results,
        "summary": build_summary(results),
    }
    report_path.write_text(json.dumps(report, indent=2, ensure_ascii=True), encoding="utf-8")
    summary_path.write_text(build_markdown_summary(report), encoding="utf-8")

    print(
        json.dumps(
            {
                "report_dir": str(report_dir),
                "report_path": str(report_path),
                "summary_path": str(summary_path),
                "passed": report["summary"]["passed"],
                "failed_assertions": report["summary"]["failed_assertions"],
                "generation_failed": report["summary"]["generation_failed"],
                "errors": report["summary"]["errors"],
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))
