from __future__ import annotations

import argparse
import asyncio
import json
import shutil
from dataclasses import asdict, dataclass, field
from pathlib import Path
from uuid import uuid4

from ai_orchestrator.config.settings import PROJECT_ROOT, Settings
from ai_orchestrator.schemas.task_request import TaskConstraints, TaskRequest
from ai_orchestrator.services.task_service import build_task_service, build_task_worker


@dataclass(slots=True)
class EvalCase:
    case_id: str
    description: str
    goal: str
    repo_dir: Path
    labels: list[str]
    constraints: dict
    expect: dict


@dataclass(slots=True)
class EvalCaseResult:
    case_id: str
    description: str
    labels: list[str]
    passed: bool
    status: str
    expected_status: str
    step_statuses: dict[str, str] = field(default_factory=dict)
    focus_issue_types: list[str] = field(default_factory=list)
    failed_checks: list[str] = field(default_factory=list)
    issues: list[str] = field(default_factory=list)


@dataclass(slots=True)
class EvalRunSummary:
    passed: int
    failed: int
    total: int
    results: list[EvalCaseResult] = field(default_factory=list)


def load_eval_cases(cases_dir: Path, case_names: list[str] | None = None) -> list[EvalCase]:
    selected = set(case_names or [])
    cases: list[EvalCase] = []
    for case_dir in sorted(path for path in cases_dir.iterdir() if path.is_dir()):
        if selected and case_dir.name not in selected:
            continue
        payload = json.loads((case_dir / "task.json").read_text(encoding="utf-8"))
        cases.append(
            EvalCase(
                case_id=case_dir.name,
                description=str(payload.get("description", case_dir.name)),
                goal=str(payload["goal"]),
                repo_dir=case_dir / "repo",
                labels=[str(label) for label in payload.get("labels", [])],
                constraints=dict(payload.get("constraints", {})),
                expect=dict(payload.get("expect", {})),
            )
        )
    return cases


async def _run_case(case: EvalCase, base_settings: Settings, use_mock_models: bool) -> EvalCaseResult:
    eval_temp_root = PROJECT_ROOT / "artifacts" / "eval-runs"
    eval_temp_root.mkdir(parents=True, exist_ok=True)
    temp_dir = eval_temp_root / f"{case.case_id}-{uuid4().hex[:8]}"
    temp_dir.mkdir(parents=True, exist_ok=True)
    try:
        repo_copy = temp_dir / "repo"
        shutil.copytree(case.repo_dir, repo_copy)
        model_config_path = base_settings.model_config_path
        if use_mock_models:
            model_config_path = temp_dir / "models.yaml"
            model_config_path.write_text(
                "\n".join(
                    [
                        "steps:",
                        "  planning:",
                        "    model: mock/planner",
                        "    temperature: 0.1",
                        "    max_tokens: 800",
                        "  coding:",
                        "    model: mock/coder",
                        "    temperature: 0.0",
                        "    max_tokens: 2400",
                        "  analysis:",
                        "    model: mock/reasoner",
                        "    temperature: 0.1",
                        "    max_tokens: 1200",
                        "  failure_analysis:",
                        "    model: mock/reasoner",
                        "    temperature: 0.1",
                        "    max_tokens: 1200",
                        "  summary:",
                        "    model: mock/summary",
                        "    temperature: 0.1",
                        "    max_tokens: 800",
                        "  embeddings:",
                        "    model: mock/embeddings",
                        "    temperature: 0.0",
                        "    max_tokens: 0",
                    ]
                )
                + "\n",
                encoding="utf-8",
                newline="\n",
            )

        settings = Settings(
            _env_file=None,
            queue_backend="sqlite",
            model_config_path=model_config_path,
            data_dir=temp_dir / "data",
            artifacts_dir=temp_dir / "artifacts",
            database_path=temp_dir / "data" / "tasks.db",
            memory_database_path=temp_dir / "data" / "memory.db",
            workspaces_dir=temp_dir / "workspaces",
        )
        service = build_task_service(settings)
        worker = build_task_worker(settings)
        try:
            request = TaskRequest(
                goal=case.goal,
                repo_path=repo_copy,
                constraints=TaskConstraints.model_validate(case.constraints),
            )
            submitted = await service.submit_task(request)
            worker_handle = asyncio.create_task(worker.run_until_idle(idle_cycles=5))
            task = await service.wait_for_completion(submitted.task_id, timeout_seconds=20.0)
            await worker_handle
        finally:
            await service.shutdown()

        issues: list[str] = []
        expected_status = str(case.expect.get("status", "completed"))
        if task.status.value != expected_status:
            issues.append(f"Expected status `{expected_status}`, got `{task.status.value}`.")

        step_statuses: dict[str, str] = {}
        for step in task.step_history:
            step_statuses[step.step] = step.status
        for step_name, expected_step_status in case.expect.get("step_statuses", {}).items():
            actual_step_status = step_statuses.get(step_name, "missing")
            if actual_step_status != expected_step_status:
                issues.append(
                    f"Expected step `{step_name}` to be `{expected_step_status}`, got `{actual_step_status}`."
                )

        for relative_path, required_fragments in case.expect.get("file_contains", {}).items():
            content = (repo_copy / relative_path).read_text(encoding="utf-8")
            for fragment in required_fragments:
                if fragment not in content:
                    issues.append(f"`{relative_path}` is missing required fragment: {fragment}")

        for relative_path, forbidden_fragments in case.expect.get("file_not_contains", {}).items():
            content = (repo_copy / relative_path).read_text(encoding="utf-8")
            for fragment in forbidden_fragments:
                if fragment in content:
                    issues.append(f"`{relative_path}` still contains forbidden fragment: {fragment}")

        focus_issue_types = sorted(
            {
                issue.issue_type
                for issue in (task.focus_issue_inventory.issues if task.focus_issue_inventory is not None else [])
            }
        )
        for issue_type in case.expect.get("focus_issue_types_present", []):
            if issue_type not in focus_issue_types:
                issues.append(f"Expected focus issue type `{issue_type}` to be present.")
        for issue_type in case.expect.get("focus_issue_types_absent", []):
            if issue_type in focus_issue_types:
                issues.append(f"Expected focus issue type `{issue_type}` to be absent.")

        failed_checks = task.validation_history[-1].failed_checks if task.validation_history else []
        expected_failed_checks = case.expect.get("failed_checks")
        if expected_failed_checks is not None:
            if sorted(expected_failed_checks) != sorted(failed_checks):
                issues.append(
                    f"Expected failed checks {sorted(expected_failed_checks)}, got {sorted(failed_checks)}."
                )

        error_contains = str(case.expect.get("error_contains", "")).strip()
        if error_contains and error_contains not in (task.error_message or ""):
            issues.append(f"Expected error message to contain: {error_contains}")

        return EvalCaseResult(
            case_id=case.case_id,
            description=case.description,
            labels=case.labels,
            passed=not issues,
            status=task.status.value,
            expected_status=expected_status,
            step_statuses=step_statuses,
            focus_issue_types=focus_issue_types,
            failed_checks=failed_checks,
            issues=issues,
        )
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


async def run_eval_cases(
    *,
    case_names: list[str] | None = None,
    cases_dir: Path | None = None,
    use_mock_models: bool = False,
) -> EvalRunSummary:
    base_settings = Settings()
    root = cases_dir or (PROJECT_ROOT / "evals" / "cases")
    cases = load_eval_cases(root, case_names)
    results: list[EvalCaseResult] = []
    for case in cases:
        results.append(await _run_case(case, base_settings, use_mock_models))
    passed = sum(1 for result in results if result.passed)
    failed = len(results) - passed
    return EvalRunSummary(passed=passed, failed=failed, total=len(results), results=results)


def main() -> None:
    parser = argparse.ArgumentParser(description="Run AIIQ eval cases.")
    parser.add_argument("--cases", nargs="*", default=None, help="Optional case names to run.")
    parser.add_argument("--cases-dir", default=None, help="Optional override for eval case directory.")
    parser.add_argument("--mock-models", action="store_true", help="Use mock model config for deterministic runs.")
    parser.add_argument("--output", default=None, help="Optional JSON output path.")
    args = parser.parse_args()

    summary = asyncio.run(
        run_eval_cases(
            case_names=args.cases,
            cases_dir=Path(args.cases_dir) if args.cases_dir else None,
            use_mock_models=args.mock_models,
        )
    )
    payload = asdict(summary)
    text = json.dumps(payload, indent=2)
    if args.output:
        Path(args.output).write_text(text + "\n", encoding="utf-8")
    print(text)


if __name__ == "__main__":
    main()
