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

import argparse
import asyncio
import json
import os
import sys
import time
from collections import Counter, defaultdict
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 / "quality_cases.json"
DEFAULT_REPORTS_ROOT = HERE / "quality_reports"


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("quality 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 filter_engines(
    engines: list[dict[str, Any]],
    selected_engine_ids: set[str],
    max_engines: int | None,
) -> list[dict[str, Any]]:
    wrappers = [engine for engine in engines if engine.get("can_be_wrapper")]
    if selected_engine_ids:
        wrappers = [engine for engine in wrappers if str(engine.get("engine_id")) in selected_engine_ids]
    wrappers = sorted(wrappers, key=lambda engine: str(engine.get("engine_id")))
    if max_engines:
        wrappers = wrappers[: max_engines]
    return wrappers


async def run_case_for_engine(
    case: dict[str, Any],
    engine: dict[str, Any],
) -> dict[str, Any]:
    engine_id = str(engine.get("engine_id"))
    started = time.perf_counter()
    started_at = datetime.now(timezone.utc).isoformat()
    telemetry_context = {
        "source_surface": "quality_harness",
        "source_mode": "engine_coverage",
        "quality_case_id": case.get("id"),
        "quality_case_title": case.get("title"),
        "quality_wrapper_engine_id": engine_id,
    }

    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=engine_id,
            section_count=case.get("section_count"),
            uploaded_content=case.get("uploaded_content"),
            interaction_mode=case.get("interaction_mode"),
            telemetry_context=telemetry_context,
            ctx=None,
        )
        duration_ms = int((time.perf_counter() - started) * 1000)
        artifact_id = response.get("artifact_id")
        state = mcp_server.load_artifact_state(artifact_id) if artifact_id else {}
        current = state.get("current") if isinstance(state, dict) else {}
        local_path = current.get("local_path")
        local_uri = Path(local_path).resolve().as_uri() if local_path else None
        return {
            "status": "succeeded",
            "started_at": started_at,
            "duration_ms": duration_ms,
            "case_id": case.get("id"),
            "case_title": case.get("title"),
            "baseline_url": case.get("baseline_url"),
            "engine_id": engine_id,
            "engine_name": engine.get("name"),
            "artifact_id": artifact_id,
            "title": response.get("title"),
            "document_format": response.get("document_format"),
            "summary": response.get("summary"),
            "quality_score": response.get("quality_score"),
            "warnings": response.get("warnings") or [],
            "warning_types": sorted(
                {
                    str(warning.get("type"))
                    for warning in (response.get("warnings") or [])
                    if isinstance(warning, dict) and warning.get("type")
                }
            ),
            "preview_url": response.get("preview_url"),
            "artifact_page_url": response.get("preview_url") or local_uri,
            "local_path": local_path,
            "local_uri": local_uri,
            "wrapper_engine_id": response.get("wrapper_engine_id"),
            "section_count": response.get("section_count"),
            "section_titles": response.get("section_titles") or [],
            "interaction_mode": response.get("interaction_mode"),
            "structured_data_present": bool(case.get("data")),
            "notes": response.get("notes") or [],
        }
    except Exception as exc:  # noqa: BLE001
        duration_ms = int((time.perf_counter() - started) * 1000)
        return {
            "status": "failed",
            "started_at": started_at,
            "duration_ms": duration_ms,
            "case_id": case.get("id"),
            "case_title": case.get("title"),
            "baseline_url": case.get("baseline_url"),
            "engine_id": engine_id,
            "engine_name": engine.get("name"),
            "error": str(exc),
        }


def build_summary(results: list[dict[str, Any]]) -> dict[str, Any]:
    succeeded = [result for result in results if result.get("status") == "succeeded"]
    failed = [result for result in results if result.get("status") == "failed"]
    warning_counter: Counter[str] = Counter()
    engine_rollup: dict[str, dict[str, Any]] = {}
    case_rollup: dict[str, dict[str, Any]] = {}

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

    engine_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
    case_groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for result in results:
        engine_groups[str(result.get("engine_id"))].append(result)
        case_groups[str(result.get("case_id"))].append(result)

    for engine_id, group in engine_groups.items():
        engine_successes = [item for item in group if item.get("status") == "succeeded"]
        engine_rollup[engine_id] = {
            "engine_name": next((item.get("engine_name") for item in group if item.get("engine_name")), engine_id),
            "runs": len(group),
            "successes": len(engine_successes),
            "failures": len(group) - len(engine_successes),
            "avg_quality_score": round(
                sum(int(item.get("quality_score") or 0) for item in engine_successes) / len(engine_successes),
                1,
            )
            if engine_successes
            else None,
            "warning_types": sorted(
                {
                    warning_type
                    for item in engine_successes
                    for warning_type in (item.get("warning_types") or [])
                }
            ),
        }

    for case_id, group in case_groups.items():
        case_successes = [item for item in group if item.get("status") == "succeeded"]
        case_rollup[case_id] = {
            "case_title": next((item.get("case_title") for item in group if item.get("case_title")), case_id),
            "runs": len(group),
            "successes": len(case_successes),
            "failures": len(group) - len(case_successes),
            "avg_quality_score": round(
                sum(int(item.get("quality_score") or 0) for item in case_successes) / len(case_successes),
                1,
            )
            if case_successes
            else None,
        }

    return {
        "total_runs": len(results),
        "succeeded": len(succeeded),
        "failed": len(failed),
        "avg_quality_score": round(
            sum(int(item.get("quality_score") or 0) for item in succeeded) / len(succeeded),
            1,
        )
        if succeeded
        else None,
        "warning_types": dict(sorted(warning_counter.items())),
        "by_engine": dict(sorted(engine_rollup.items())),
        "by_case": dict(sorted(case_rollup.items())),
    }


def build_markdown_summary(report: dict[str, Any]) -> str:
    summary = report["summary"]
    lines = [
        "# IterFact MCP Quality Harness",
        "",
        f"- Run at: `{report['generated_at']}`",
        f"- Total runs: `{summary['total_runs']}`",
        f"- Succeeded: `{summary['succeeded']}`",
        f"- Failed: `{summary['failed']}`",
        f"- Avg quality score: `{summary['avg_quality_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("## Engine rollup")
    for engine_id, engine_summary in summary["by_engine"].items():
        lines.append(
            f"- `{engine_id}` {engine_summary['engine_name']}: "
            f"{engine_summary['successes']}/{engine_summary['runs']} passed, "
            f"avg quality `{engine_summary['avg_quality_score']}`"
        )
    lines.append("")
    lines.append("## Case rollup")
    for case_id, case_summary in summary["by_case"].items():
        lines.append(
            f"- `{case_id}` {case_summary['case_title']}: "
            f"{case_summary['successes']}/{case_summary['runs']} passed, "
            f"avg quality `{case_summary['avg_quality_score']}`"
        )
    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 IterFact MCP quality coverage across wrapper engines.")
    parser.add_argument("--cases", default=str(DEFAULT_CASES_PATH), help="Path to the quality cases JSON file.")
    parser.add_argument("--case", action="append", dest="case_ids", default=[], help="Case id to include. Repeatable.")
    parser.add_argument("--engine", action="append", dest="engine_ids", default=[], help="Wrapper engine id to include. Repeatable.")
    parser.add_argument("--max-cases", type=int, default=None, help="Limit the number of cases.")
    parser.add_argument("--max-engines", type=int, default=None, help="Limit the number of wrapper engines.")
    parser.add_argument("--output-root", default=str(DEFAULT_REPORTS_ROOT), help="Directory where report folders should be created.")
    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 quality cases matched the requested filters.")

    engine_payload = await mcp_server.list_engines()
    engines = filter_engines(engine_payload.get("engines") or [], set(args.engine_ids), args.max_engines)
    if not engines:
        raise ValueError("No wrapper engines 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:
        case_engine_ids = set(str(engine_id) for engine_id in (case.get("engines") or []) if engine_id)
        case_engines = [engine for engine in engines if not case_engine_ids or str(engine.get("engine_id")) in case_engine_ids]
        for engine in case_engines:
            results.append(await run_case_for_engine(case, engine))

    report_path = report_dir / "coverage_report.json"
    summary_path = report_dir / "coverage_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),
        "engine_count": len(engines),
        "cases": cases,
        "engines": engines,
        "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),
                "succeeded": report["summary"]["succeeded"],
                "failed": report["summary"]["failed"],
            },
            indent=2,
        )
    )
    return 0


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