#!/usr/bin/env python3
"""Select an execution target from a runtime-injected catalog.

The selector performs no quota lookup.  Every initial candidate is eligible;
runtime failures such as ``provider-quota`` advance to the next catalog entry.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path


SCHEMA_VERSION = "2.0"
CATALOG_ENV = "AGENT_TASK_EXECUTION_CATALOG"
TIMEZONE_NAME = "UTC"
_FILENAME_RE = re.compile(r"^(PLAN|CODE_REVIEW)-(local|cloud)-G(\d{2})\.md$")
_MILESTONE_TASK_ID_PATTERN = r"[A-Za-z0-9]+(?:[-_+=][A-Za-z0-9]+){0,3}"
_MILESTONE_TASK_ID_RE = re.compile(rf"\A{_MILESTONE_TASK_ID_PATTERN}\Z")
_HEADER_RE = re.compile(
    r"\A<!--\s*task=(?P<task>\S+)\s+plan=(?P<plan>\d+)\s+tag=(?P<tag>\S+)"
    r"(?:\s+milestone-task=(?P<milestone_task>[^,\s]+(?:,[^,\s]+)*))?"
    r"\s*-->[ \t]*(?:\r?\n|\Z)"
)
_STAGE_BY_KIND = {"PLAN": "worker", "CODE_REVIEW": "review"}
_VALID_TRANSITIONS = {"initial", "resume", "failover"}
_QUALIFIED_FAILOVER_FAILURES = {
    "provider-quota",
    "context-limit",
    "model-unavailable",
    "provider-stream-disconnect",
    "provider-connection",
}


def _load_policy():
    path = Path(__file__).resolve().parent / "execution_target_policy.py"
    spec = importlib.util.spec_from_file_location("execution_target_policy", path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"failed to load execution target policy: {path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


policy = _load_policy()


class SelectorInputError(Exception):
    """Input contract violation returned as stderr JSON with a non-zero exit."""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code


def resolve_catalog_path(value: str | Path | None = None) -> Path:
    raw = str(value) if value is not None else os.environ.get(CATALOG_ENV, "")
    if not raw:
        raise SelectorInputError(
            "missing_execution_catalog",
            f"inject the execution catalog with --catalog or {CATALOG_ENV}",
        )
    return Path(raw).expanduser().resolve()


def load_runtime_catalog(value: str | Path | None = None):
    try:
        return policy.load_catalog(resolve_catalog_path(value))
    except SelectorInputError:
        raise
    except (OSError, ValueError) as exc:
        raise SelectorInputError("invalid_execution_catalog", str(exc)) from exc


def _parse_filename(task_file: Path) -> tuple[str, str, int]:
    name = Path(task_file).name
    match = _FILENAME_RE.match(name)
    if match is None:
        raise SelectorInputError(
            "invalid_task_filename",
            f"task file must match (PLAN|CODE_REVIEW)-(local|cloud)-GNN.md: {name!r}",
        )
    kind, lane, grade_str = match.group(1), match.group(2), match.group(3)
    grade = int(grade_str)
    if not 1 <= grade <= 10:
        raise SelectorInputError("invalid_grade", f"grade must be G01..G10: G{grade_str}")
    return kind, lane, grade


def _parse_header(task_file: Path) -> tuple[str, int, str, str | None]:
    try:
        with Path(task_file).open("rb") as handle:
            text = handle.read(1024).decode("utf-8", errors="replace")
    except OSError as exc:
        raise SelectorInputError("task_file_unreadable", str(exc)) from exc
    match = _HEADER_RE.search(text)
    if match is None:
        raise SelectorInputError(
            "malformed_header",
            "first line must contain <!-- task=... plan=N tag=... "
            "[milestone-task=id[,id...]] -->",
        )
    task = match.group("task")
    milestone_task = match.group("milestone_task")
    task_ids = tuple(milestone_task.split(",")) if milestone_task else ()
    invalid_ids = [item for item in task_ids if _MILESTONE_TASK_ID_RE.fullmatch(item) is None]
    if invalid_ids:
        raise SelectorInputError(
            "invalid_milestone_task",
            "milestone-task ids must follow the Milestone item-id grammar: "
            + ", ".join(invalid_ids),
        )
    if len(task_ids) != len(set(task_ids)):
        raise SelectorInputError(
            "duplicate_milestone_task",
            "milestone-task must contain unique comma-separated Task ids",
        )
    milestone_group = task.split("/", 1)[0].startswith("m-")
    if milestone_group and not milestone_task:
        raise SelectorInputError(
            "missing_milestone_task",
            "m-* task headers require milestone-task=id[,id...]",
        )
    if not milestone_group and milestone_task:
        raise SelectorInputError(
            "unexpected_milestone_task",
            "non-milestone task headers must omit milestone-task",
        )
    return task, int(match.group("plan")), match.group("tag"), milestone_task


def _work_unit_id(header: tuple[str, int, str, str | None]) -> str:
    task, plan, tag, milestone_task = header
    result = f"{task}::plan-{plan}::tag-{tag}"
    if milestone_task:
        result += f"::milestone-task-{milestone_task}"
    return result


def _target_snapshot(target) -> dict:
    return {
        "target_id": target.catalog_id,
        "agent": target.agent,
        "model": target.model,
        "execution_class": target.execution_class,
        "selfcheck_required": target.selfcheck_required,
    }


def _candidate_snapshot(target, rank: int) -> dict:
    return {"candidate_rank": rank, **_target_snapshot(target)}


def _validate_target_snapshot(value: object, prefix: str) -> dict:
    code = "malformed_prior_decision"
    if not isinstance(value, dict):
        raise SelectorInputError(code, f"{prefix} must be an object")
    required = {
        "target_id",
        "agent",
        "model",
        "execution_class",
        "selfcheck_required",
    }
    missing = required - set(value)
    if missing:
        raise SelectorInputError(code, f"{prefix} missing keys: {sorted(missing)}")
    for field in ("target_id", "agent", "model"):
        if not isinstance(value[field], str) or not value[field]:
            raise SelectorInputError(code, f"{prefix}.{field} must be a non-empty string")
    if value["execution_class"] not in policy.VALID_EXECUTION_CLASSES:
        raise SelectorInputError(
            code,
            f"{prefix}.execution_class must be one of {sorted(policy.VALID_EXECUTION_CLASSES)}",
        )
    if not isinstance(value["selfcheck_required"], bool):
        raise SelectorInputError(code, f"{prefix}.selfcheck_required must be a boolean")
    return value


def _validate_prior_decision(value: object) -> dict:
    code = "malformed_prior_decision"
    if not isinstance(value, dict):
        raise SelectorInputError(code, "prior_decision must be an object")
    required = {
        "schema_version",
        "work_unit_id",
        "stage",
        "lane",
        "grade",
        "catalog",
        "selected",
        "candidates",
        "decision",
        "transition",
    }
    missing = required - set(value)
    if missing:
        raise SelectorInputError(code, f"prior_decision missing keys: {sorted(missing)}")
    if value["schema_version"] != SCHEMA_VERSION:
        raise SelectorInputError(code, f"prior_decision.schema_version must be {SCHEMA_VERSION!r}")
    if value["stage"] not in policy.VALID_STAGES or value["lane"] not in policy.VALID_LANES:
        raise SelectorInputError(code, "prior_decision stage/lane is invalid")
    grade = value["grade"]
    if isinstance(grade, bool) or not isinstance(grade, int) or not 1 <= grade <= 10:
        raise SelectorInputError(code, "prior_decision.grade must be G01..G10")
    if not isinstance(value["work_unit_id"], str) or not value["work_unit_id"]:
        raise SelectorInputError(code, "prior_decision.work_unit_id must be a non-empty string")
    catalog = value["catalog"]
    if not isinstance(catalog, dict):
        raise SelectorInputError(code, "prior_decision.catalog must be an object")
    for field in ("schema_version", "revision", "source", "route_id"):
        if not isinstance(catalog.get(field), str) or not catalog[field]:
            raise SelectorInputError(code, f"prior_decision.catalog.{field} must be a non-empty string")
    _validate_target_snapshot(value["selected"], "prior_decision.selected")
    candidates = value["candidates"]
    if not isinstance(candidates, list) or not candidates:
        raise SelectorInputError(code, "prior_decision.candidates must be a non-empty list")
    for index, candidate in enumerate(candidates, 1):
        _validate_target_snapshot(candidate, f"prior_decision.candidates[{index - 1}]")
        if candidate.get("candidate_rank") != index:
            raise SelectorInputError(
                code,
                f"prior_decision.candidates[{index - 1}].candidate_rank must be {index}",
            )
    decision = value["decision"]
    if not isinstance(decision, dict):
        raise SelectorInputError(code, "prior_decision.decision must be an object")
    for field in ("rule_id", "evaluated_at", "timezone", "time_window"):
        if not isinstance(decision.get(field), str) or not decision[field]:
            raise SelectorInputError(code, f"prior_decision.decision.{field} must be a non-empty string")
    if not isinstance(decision.get("policy_priority"), int) or isinstance(
        decision["policy_priority"], bool
    ):
        raise SelectorInputError(code, "prior_decision.decision.policy_priority must be an integer")
    if not isinstance(decision.get("reason_codes"), list) or not all(
        isinstance(item, str) and item for item in decision["reason_codes"]
    ):
        raise SelectorInputError(code, "prior_decision.decision.reason_codes must be a string list")
    if not isinstance(decision.get("pinned"), bool):
        raise SelectorInputError(code, "prior_decision.decision.pinned must be a boolean")
    transition = value["transition"]
    if not isinstance(transition, dict) or transition.get("trigger") not in _VALID_TRANSITIONS:
        raise SelectorInputError(code, "prior_decision.transition is invalid")
    return value


def _catalog_matches_prior(catalog, prior: dict, decision) -> None:
    code = "catalog_revision_mismatch"
    evidence = prior["catalog"]
    if evidence["schema_version"] != policy.CATALOG_SCHEMA_VERSION:
        raise SelectorInputError(code, "persisted catalog schema is unsupported")
    if evidence["revision"] != catalog.revision:
        raise SelectorInputError(
            code,
            "the injected execution catalog changed after this work unit was selected",
        )
    if evidence["route_id"] != decision.route_id:
        raise SelectorInputError(code, "persisted catalog route does not match the task route")


def _validate_prior_candidate_identity(prior: dict, *, catalog, decision) -> None:
    code = "malformed_prior_decision"
    expected = [_candidate_snapshot(item, rank) for rank, item in enumerate(decision.candidates, 1)]
    if prior["candidates"] != expected:
        raise SelectorInputError(code, "prior_decision candidates do not match the injected catalog route")
    selected = prior["selected"]
    if selected not in [{key: value for key, value in item.items() if key != "candidate_rank"} for item in expected]:
        raise SelectorInputError(code, "prior_decision selected target is not in the injected route")


def _base_decision(
    *,
    catalog,
    route,
    work_unit_id: str,
    stage: str,
    lane: str,
    grade: int,
    evaluated_at: datetime,
    selected,
    pinned: bool,
    previous_target: dict | None,
    trigger: str,
) -> dict:
    selected_snapshot = _target_snapshot(selected)
    return {
        "schema_version": SCHEMA_VERSION,
        "work_unit_id": work_unit_id,
        "stage": stage,
        "lane": lane,
        "grade": grade,
        "catalog": {
            "schema_version": policy.CATALOG_SCHEMA_VERSION,
            "revision": catalog.revision,
            "source": str(catalog.source),
            "route_id": route.route_id,
        },
        "selected": selected_snapshot,
        "candidates": [
            _candidate_snapshot(item, rank)
            for rank, item in enumerate(route.candidates, 1)
        ],
        "decision": {
            "rule_id": route.rule_id,
            "policy_priority": route.policy_priority,
            "reason_codes": list(route.reason_codes),
            "evaluated_at": evaluated_at.astimezone(timezone.utc).isoformat(),
            "timezone": TIMEZONE_NAME,
            "time_window": route.time_window,
            "pinned": pinned,
        },
        "transition": {
            "previous_target": previous_target,
            "next_target": selected_snapshot,
            "trigger": trigger,
            "context_transfer": "logical" if trigger == "failover" else "none",
        },
    }


def select_execution_target_for_route(
    *,
    work_unit_id: str,
    stage: str,
    lane: str,
    grade: int,
    evaluated_at: datetime,
    catalog_path: str | Path | None = None,
    transition: str = "initial",
    prior_decision: dict | None = None,
    failure_class: str | None = None,
) -> dict:
    if transition not in _VALID_TRANSITIONS:
        raise SelectorInputError("invalid_transition", f"unsupported transition: {transition}")
    if evaluated_at.tzinfo is None or evaluated_at.utcoffset() is None:
        raise SelectorInputError("naive_evaluated_at", "evaluated_at must be timezone-aware")
    catalog = load_runtime_catalog(catalog_path)
    try:
        route = policy.select_policy(
            catalog=catalog,
            stage=stage,
            lane=lane,
            grade=grade,
            evaluated_at=evaluated_at,
        )
    except ValueError as exc:
        raise SelectorInputError("invalid_route", str(exc)) from exc
    if transition == "initial":
        if prior_decision is not None:
            raise SelectorInputError("unexpected_prior_decision", "initial transition must not include prior_decision")
        return _base_decision(
            catalog=catalog,
            route=route,
            work_unit_id=work_unit_id,
            stage=stage,
            lane=lane,
            grade=grade,
            evaluated_at=evaluated_at,
            selected=route.candidates[0],
            pinned=False,
            previous_target=None,
            trigger="initial",
        )
    prior = _validate_prior_decision(prior_decision)
    expected_identity = (work_unit_id, stage, lane, grade)
    actual_identity = (
        prior["work_unit_id"],
        prior["stage"],
        prior["lane"],
        prior["grade"],
    )
    if actual_identity != expected_identity:
        raise SelectorInputError("prior_decision_mismatch", "prior_decision belongs to a different work unit or route")
    _catalog_matches_prior(catalog, prior, route)
    _validate_prior_candidate_identity(prior, catalog=catalog, decision=route)
    selected_id = prior["selected"]["target_id"]
    selected_index = [item.catalog_id for item in route.candidates].index(selected_id)
    if transition == "resume":
        selected = route.candidates[selected_index]
        return _base_decision(
            catalog=catalog,
            route=route,
            work_unit_id=work_unit_id,
            stage=stage,
            lane=lane,
            grade=grade,
            evaluated_at=evaluated_at,
            selected=selected,
            pinned=True,
            previous_target=_target_snapshot(selected),
            trigger="resume",
        )
    if failure_class not in _QUALIFIED_FAILOVER_FAILURES:
        raise SelectorInputError(
            "unqualified_failover",
            f"failure_class does not qualify for target failover: {failure_class!r}",
        )
    next_index = selected_index + 1
    if next_index >= len(route.candidates):
        raise SelectorInputError("no_failover_candidate", "the injected route has no unused next target")
    return _base_decision(
        catalog=catalog,
        route=route,
        work_unit_id=work_unit_id,
        stage=stage,
        lane=lane,
        grade=grade,
        evaluated_at=evaluated_at,
        selected=route.candidates[next_index],
        pinned=False,
        previous_target=dict(prior["selected"]),
        trigger="failover",
    )


def select_execution_target(
    task_file: Path,
    *,
    stage: str | None = None,
    evaluated_at: datetime | None = None,
    catalog_path: str | Path | None = None,
    transition: str = "initial",
    prior_decision: dict | None = None,
    failure_class: str | None = None,
) -> dict:
    kind, lane, grade = _parse_filename(Path(task_file))
    inferred_stage = _STAGE_BY_KIND[kind]
    if stage is not None and stage != inferred_stage:
        raise SelectorInputError(
            "stage_mismatch",
            f"stage {stage!r} does not match task filename stage {inferred_stage!r}",
        )
    return select_execution_target_for_route(
        work_unit_id=_work_unit_id(_parse_header(Path(task_file))),
        stage=inferred_stage,
        lane=lane,
        grade=grade,
        evaluated_at=evaluated_at or datetime.now(timezone.utc),
        catalog_path=catalog_path,
        transition=transition,
        prior_decision=prior_decision,
        failure_class=failure_class,
    )


def to_json(payload: dict) -> str:
    return json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def _load_json_arg(value: str | None):
    if value is None:
        return None
    path = Path(value)
    try:
        return json.loads(path.read_text(encoding="utf-8")) if path.is_file() else json.loads(value)
    except (OSError, json.JSONDecodeError) as exc:
        raise SelectorInputError("invalid_json_argument", str(exc)) from exc


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("task_file", type=Path)
    parser.add_argument("--stage", choices=sorted(policy.VALID_STAGES))
    parser.add_argument("--catalog")
    parser.add_argument("--evaluated-at")
    parser.add_argument("--transition", choices=sorted(_VALID_TRANSITIONS), default="initial")
    parser.add_argument("--prior-decision")
    parser.add_argument("--failure-class")
    args = parser.parse_args(argv)
    try:
        evaluated_at = datetime.fromisoformat(args.evaluated_at) if args.evaluated_at else None
        payload = select_execution_target(
            args.task_file,
            stage=args.stage,
            evaluated_at=evaluated_at,
            catalog_path=args.catalog,
            transition=args.transition,
            prior_decision=_load_json_arg(args.prior_decision),
            failure_class=args.failure_class,
        )
    except (SelectorInputError, ValueError) as exc:
        print(
            to_json({"error": {"code": getattr(exc, "code", "invalid_input"), "message": str(exc)}}),
            file=sys.stderr,
        )
        return 2
    print(to_json(payload))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())