import argparse
import asyncio
import importlib.util
import json
import os
import stat
import subprocess
import sys
import unittest
from datetime import datetime, timezone
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest import mock


SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "dispatch.py"
SPEC = importlib.util.spec_from_file_location("agent_task_dispatch_test", SCRIPT)
dispatch = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
sys.modules[SPEC.name] = dispatch
SPEC.loader.exec_module(dispatch)


def catalog_value(command: str = "/bin/true") -> dict:
    targets = {
        "primary": {
            "agent": "runner-primary",
            "model": "model-primary",
            "execution_class": "local_model",
            "selfcheck_required": True,
            "runtime": {
                "command": [command, "--workspace", "{workspace}", "--model", "{model}", "{prompt}"],
                "resume_command": [command, "--resume", "{resume_session}", "{prompt}"],
                "environment": {"TARGET_ID": "{target_id}"},
                "output_format": "jsonl",
                "native_session_monitor": True,
                "session_path": "sessions/{session_id}.jsonl",
            },
        },
        "alternate": {
            "agent": "runner-alternate",
            "model": "model-alternate",
            "execution_class": "cloud_model",
            "runtime": {"command": [command, "{prompt}"]},
        },
    }
    routes = {"worker": {}, "review": {}}
    for stage in routes:
        for lane in ("local", "cloud"):
            for grade in range(1, 11):
                routes[stage][f"{lane}-G{grade:02d}"] = {
                    "candidates": ["primary", "alternate"],
                    "rule_id": f"{stage}-{lane}-{grade:02d}",
                    "reason_codes": ["injected-route"],
                }
    return {"schema_version": "1.0", "targets": targets, "routes": routes}


def write_catalog(root: Path, value: dict | None = None) -> Path:
    path = root / "execution-catalog.json"
    path.write_text(json.dumps(value or catalog_value()), encoding="utf-8")
    return path


def write_plan(root: Path, *, task_name: str = "group/01_task") -> Path:
    directory = root / "agent-task" / task_name
    directory.mkdir(parents=True, exist_ok=True)
    path = directory / "PLAN-cloud-G05.md"
    path.write_text(
        f"<!-- task={task_name} plan=0 tag=API -->\n\n"
        "# Plan\n\n## Modified Files Summary\n\n"
        "| File | Action |\n|---|---|\n| `src/item.txt` | modify |\n",
        encoding="utf-8",
    )
    return path


def task_from_plan(root: Path, plan: Path) -> dispatch.Task:
    directory = plan.parent
    return dispatch.Task(
        name="group/01_task",
        directory=directory,
        plan=plan,
        review=None,
        user_review=None,
        recovery=False,
        index=1,
        write_set={"src/item.txt"},
        write_set_known=True,
        plan_hash=dispatch.sha256_file(plan),
    )


class RuntimeCatalogDispatcherTests(unittest.TestCase):
    def setUp(self):
        self.previous_catalog = dispatch.EXECUTION_CATALOG_PATH

    def tearDown(self):
        dispatch.EXECUTION_CATALOG_PATH = self.previous_catalog

    def test_agent_spec_is_loaded_from_persisted_catalog_evidence(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            catalog = write_catalog(root)
            plan = write_plan(root)
            dispatch.EXECUTION_CATALOG_PATH = catalog
            selector = dispatch._selector_module()
            decision = selector.select_execution_target(plan, catalog_path=catalog)
            spec = dispatch.agent_spec_from_decision(decision)
        self.assertEqual(spec.target_id, "primary")
        self.assertEqual(spec.cli, "runner-primary")
        self.assertEqual(spec.model, "model-primary")
        self.assertTrue(spec.native_resume)
        self.assertEqual(spec.runtime["command"][0], "/bin/true")

    def test_agent_spec_rejects_catalog_change_after_selection(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            catalog = write_catalog(root)
            plan = write_plan(root)
            selector = dispatch._selector_module()
            decision = selector.select_execution_target(plan, catalog_path=catalog)
            changed = catalog_value()
            changed["targets"]["primary"]["model"] = "changed"
            catalog.write_text(json.dumps(changed), encoding="utf-8")
            with self.assertRaisesRegex(dispatch.ExecutionDecisionError, "변경"):
                dispatch.agent_spec_from_decision(decision)

    def test_command_is_expanded_only_from_runtime_template(self):
        spec = dispatch.AgentSpec(
            "opaque-agent",
            "opaque-model",
            "opaque-agent/opaque-model",
            target_id="opaque-id",
            runtime={
                "command": ["runner", "{workspace}", "{model}", "{session_id}", "{attempt_dir}", "{prompt}"],
                "resume_command": ["runner", "resume", "{resume_session}", "{prompt}"],
            },
        )
        command = dispatch.build_command(
            spec,
            "do work",
            Path("/workspace"),
            "session-1",
            Path("/attempt"),
        )
        resumed = dispatch.build_command(
            spec,
            "continue",
            Path("/workspace"),
            "session-1",
            Path("/attempt"),
            native_resume_session=Path("/attempt/session.jsonl"),
        )
        self.assertEqual(command, ["runner", "/workspace", "opaque-model", "session-1", "/attempt", "do work"])
        self.assertEqual(resumed, ["runner", "resume", "/attempt/session.jsonl", "continue"])

    def test_preflight_checks_executable_and_optional_probe(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            value = catalog_value("/bin/true")
            value["targets"]["primary"]["runtime"]["preflight_command"] = ["/bin/true", "--check"]
            catalog = write_catalog(root, value)
            dispatch.preflight_execution_catalog(catalog)

    def test_preflight_rejects_missing_command(self):
        with TemporaryDirectory() as tmp:
            catalog = write_catalog(Path(tmp), catalog_value("definitely-missing-command"))
            with self.assertRaisesRegex(dispatch.ExecutionDecisionError, "command not found"):
                dispatch.preflight_execution_catalog(catalog)

    def test_persisted_decision_failover_uses_next_runtime_target(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            catalog = write_catalog(root)
            plan = write_plan(root)
            task = task_from_plan(root, plan)
            dispatch.EXECUTION_CATALOG_PATH = catalog
            with mock.patch.dict(os.environ, {"XDG_STATE_HOME": str(root / "state")}):
                store = dispatch.StateStore(root)
                try:
                    initial, first_spec = dispatch.persisted_execution_decision(store, task, stage="worker")
                    failed, second_spec = dispatch.persisted_execution_decision(
                        store,
                        task,
                        stage="worker",
                        transition="failover",
                        failure_class="provider-quota",
                    )
                finally:
                    store.close()
        self.assertEqual(initial["selected"]["target_id"], "primary")
        self.assertEqual(first_spec.model, "model-primary")
        self.assertEqual(failed["selected"]["target_id"], "alternate")
        self.assertEqual(second_spec.model, "model-alternate")
        self.assertNotIn("quota", failed)

    def test_retry_blocked_marks_failover_without_quota_state(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            catalog = write_catalog(root)
            plan = write_plan(root)
            task = task_from_plan(root, plan)
            dispatch.EXECUTION_CATALOG_PATH = catalog
            with mock.patch.dict(os.environ, {"XDG_STATE_HOME": str(root / "state")}):
                store = dispatch.StateStore(root)
                try:
                    decision, _ = dispatch.persisted_execution_decision(store, task, stage="worker")
                    state = store.task_state(task)
                    state.update(
                        blocked="runtime failure",
                        blocker_evidence={
                            "role": "worker",
                            "failure_class": "provider-quota",
                            "locator": "/tmp/locator.json",
                            "selected": decision["selected"],
                            "work_unit_id": decision["work_unit_id"],
                        },
                    )
                    store.save()
                    store.mark_retry_failover("group")
                    state = store.task_state(task)
                finally:
                    store.close()
        self.assertTrue(state["retry_failover_pending"])
        self.assertNotIn("quota_snapshot", state)
        self.assertNotIn("retry_quota_refresh_pending", state)

    def test_runtime_error_classifier_keeps_provider_quota(self):
        failure, evidence = dispatch.classify_failure_with_evidence(
            "HTTP 429 resource exhausted: quota reached"
        )
        self.assertEqual(failure, "provider-quota")
        self.assertIsNotNone(evidence)

    def test_generic_json_terminal_diagnostic_has_no_agent_branch(self):
        diagnostic = dispatch.terminal_diagnostic(
            "opaque-agent",
            "stdout",
            json.dumps({"type": "turn.failed", "error": {"code": 429}}),
        )
        self.assertIn("429", diagnostic or "")
        self.assertIsNone(
            dispatch.terminal_diagnostic(
                "opaque-agent",
                "stdout",
                json.dumps({"type": "message", "text": "quota design notes"}),
            )
        )

    def test_catalog_source_is_in_runtime_audit_evidence(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            catalog = write_catalog(root)
            plan = write_plan(root)
            selector = dispatch._selector_module()
            decision = selector.select_execution_target(plan, catalog_path=catalog)
            evidence = dispatch.selector_runtime_evidence(decision)
        self.assertEqual(evidence["catalog"]["source"], str(catalog.resolve()))
        self.assertNotIn("quota", evidence)


class GenericDispatcherContractTests(unittest.TestCase):
    def test_parallel_limit_contract(self):
        self.assertEqual(dispatch.validated_max_parallel(0), 0)
        self.assertEqual(dispatch.validated_max_parallel(3), 3)
        with self.assertRaises(ValueError):
            dispatch.validated_max_parallel(-1)

    def test_modified_files_summary_is_canonicalized(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            plan = write_plan(root)
            write_set, diagnostics = dispatch.inspect_write_set(plan, root)
        self.assertEqual(diagnostics, [])
        self.assertEqual(write_set, {str((root / "src/item.txt").resolve())})

    def test_outside_workspace_claim_is_rejected(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            plan = write_plan(root)
            plan.write_text(
                plan.read_text(encoding="utf-8").replace("`src/item.txt`", "`../outside.txt`"),
                encoding="utf-8",
            )
            _, diagnostics = dispatch.inspect_write_set(plan, root)
        self.assertTrue(any("outside" in item.lower() or "workspace" in item.lower() for item in diagnostics))

    def test_selector_evidence_uses_agent_model_fields(self):
        decision = {
            "work_unit_id": "group/01::plan-0::tag-API",
            "selected": {"target_id": "a", "agent": "runner", "model": "model"},
            "candidates": [
                {"candidate_rank": 1, "target_id": "a", "agent": "runner", "model": "model"}
            ],
            "decision": {"rule_id": "rule", "policy_priority": 1, "reason_codes": []},
            "transition": {"trigger": "initial"},
        }
        lines = dispatch.selector_evidence_lines(decision)
        self.assertIn("candidates=#1:runner/model", lines)
        self.assertFalse(any("quota" in line for line in lines))

    def test_validate_plan_mode_does_not_require_catalog(self):
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            plan = write_plan(root)
            completed = subprocess.run(
                [sys.executable, str(SCRIPT), "--workspace", str(root), "--validate-plan", str(plan)],
                capture_output=True,
                text=True,
                env={key: value for key, value in os.environ.items() if key != "AGENT_TASK_EXECUTION_CATALOG"},
                check=False,
            )
        self.assertEqual(completed.returncode, 0, completed.stderr)

    def test_dry_run_requires_catalog(self):
        with TemporaryDirectory() as tmp:
            completed = subprocess.run(
                [sys.executable, str(SCRIPT), "--workspace", tmp, "--dry-run"],
                capture_output=True,
                text=True,
                env={key: value for key, value in os.environ.items() if key != "AGENT_TASK_EXECUTION_CATALOG"},
                check=False,
            )
        self.assertEqual(completed.returncode, 2)
        self.assertIn("missing_execution_catalog", completed.stderr)


if __name__ == "__main__":
    unittest.main()