"""Durable, authority-bound deterministic execution plans for Online Agent Tasks."""
from __future__ import annotations

from contextlib import contextmanager
from dataclasses import dataclass
import fcntl
import hashlib
import json
from pathlib import Path
from typing import Any, Mapping

from aios_online_agent_execution import (
    GovernedExecutionBinding,
    project_execution_envelope,
)


class ExecutionPlanError(RuntimeError):
    pass


def _canonical(value: object) -> bytes:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def _digest(value: object) -> str:
    return hashlib.sha256(_canonical(value)).hexdigest()


def _task_key(task_id: str) -> str:
    readable = "".join(ch if ch.isalnum() or ch in "._-" else "-" for ch in task_id).strip("-")[:72] or "task"
    return readable + "-" + hashlib.sha256(task_id.encode("utf-8")).hexdigest()[:24]


@dataclass(frozen=True)
class ExecutionPlanRecord:
    task_id: str
    job_id: str
    source_head_oid: str
    generation: int
    idempotency_key: str
    plan_digest: str
    execution: dict[str, Any]

    def public(self) -> dict[str, Any]:
        return {
            "schema": "aios.online-agent.plan-record.v1",
            "task_id": self.task_id,
            "job_id": self.job_id,
            "source_head_oid": self.source_head_oid,
            "generation": self.generation,
            "idempotency_key": self.idempotency_key,
            "plan_digest": self.plan_digest,
            "execution": dict(self.execution),
        }


class ExecutionPlanStore:
    def __init__(self, root: Path) -> None:
        self.root = Path(root)
        self.root.mkdir(parents=True, exist_ok=True)

    def _path(self, task_id: str) -> Path:
        return self.root / (_task_key(task_id) + ".json")

    @contextmanager
    def _lock(self, task_id: str):
        path = self.root / (_task_key(task_id) + ".lock")
        with path.open("a+b") as handle:
            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
            try:
                yield
            finally:
                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)

    @staticmethod
    def _source_head(binding: GovernedExecutionBinding) -> str:
        return binding.expected_branch_oid if binding.branch_mode == "CONTINUE_EXISTING" else binding.base_oid

    def _read(self, task_id: str) -> ExecutionPlanRecord | None:
        path = self._path(task_id)
        if not path.is_file() or path.is_symlink():
            return None
        value = json.loads(path.read_text(encoding="utf-8"))
        if not isinstance(value, Mapping) or value.get("schema") != "aios.online-agent.plan-record.v1":
            raise ExecutionPlanError("EXECUTION_PLAN_RECORD_INVALID")
        execution = value.get("execution")
        if not isinstance(execution, Mapping):
            raise ExecutionPlanError("EXECUTION_PLAN_RECORD_INVALID")
        return ExecutionPlanRecord(
            task_id=str(value.get("task_id") or ""),
            job_id=str(value.get("job_id") or ""),
            source_head_oid=str(value.get("source_head_oid") or ""),
            generation=int(value.get("generation") or 0),
            idempotency_key=str(value.get("idempotency_key") or ""),
            plan_digest=str(value.get("plan_digest") or ""),
            execution=dict(execution),
        )

    def register(
        self,
        binding: GovernedExecutionBinding,
        request: Mapping[str, Any],
    ) -> dict[str, Any]:
        with self._lock(binding.task_id):
            return self._register_locked(binding, request)

    def _register_locked(
        self,
        binding: GovernedExecutionBinding,
        request: Mapping[str, Any],
    ) -> dict[str, Any]:
        if request.get("schema") != "aios.online-agent.plan.v1":
            raise ExecutionPlanError("EXECUTION_PLAN_SCHEMA_INVALID")
        if request.get("task_id") != binding.task_id or request.get("job_id") != binding.job_id:
            raise ExecutionPlanError("EXECUTION_PLAN_BINDING_CONFLICT")
        if int(request.get("generation") or 0) != binding.generation:
            raise ExecutionPlanError("EXECUTION_PLAN_GENERATION_CONFLICT")
        source_head = str(request.get("source_head_oid") or "")
        if source_head != self._source_head(binding):
            raise ExecutionPlanError("EXECUTION_PLAN_SOURCE_HEAD_CONFLICT")
        key = str(request.get("idempotency_key") or "")
        if len(key) < 8:
            raise ExecutionPlanError("EXECUTION_PLAN_IDEMPOTENCY_INVALID")
        execution_request = request.get("execution")
        if not isinstance(execution_request, Mapping) or execution_request.get("action") != "deterministic":
            raise ExecutionPlanError("EXECUTION_PLAN_DETERMINISTIC_REQUIRED")
        envelope = project_execution_envelope(binding, execution_request)
        material = {
            "task_id": binding.task_id,
            "job_id": binding.job_id,
            "source_head_oid": source_head,
            "generation": binding.generation,
            "idempotency_key": key,
            "execution": dict(execution_request),
        }
        plan_digest = _digest(material)
        existing = self._read(binding.task_id)
        expected_previous = request.get("expected_previous_plan_digest")
        if existing is not None:
            if existing.idempotency_key == key:
                if existing.plan_digest != plan_digest:
                    raise ExecutionPlanError("EXECUTION_PLAN_IDEMPOTENCY_CONFLICT")
                return {**existing.public(), "idempotent": True, "execution_digest": envelope["execution_digest"]}
            if expected_previous != existing.plan_digest:
                raise ExecutionPlanError("EXECUTION_PLAN_REPLACEMENT_CAS_MISMATCH")
            if existing.source_head_oid != source_head:
                raise ExecutionPlanError("EXECUTION_PLAN_STALE_REPLACEMENT_DENY")
        elif expected_previous not in (None, ""):
            raise ExecutionPlanError("EXECUTION_PLAN_REPLACEMENT_CAS_MISMATCH")

        record = ExecutionPlanRecord(
            task_id=binding.task_id,
            job_id=binding.job_id,
            source_head_oid=source_head,
            generation=binding.generation,
            idempotency_key=key,
            plan_digest=plan_digest,
            execution=dict(execution_request),
        )
        path = self._path(binding.task_id)
        temporary = path.with_suffix(".tmp")
        temporary.write_text(
            json.dumps(record.public(), sort_keys=True, separators=(",", ":")),
            encoding="utf-8",
        )
        temporary.replace(path)
        return {**record.public(), "idempotent": False, "execution_digest": envelope["execution_digest"]}

    def resolve(self, binding: GovernedExecutionBinding) -> dict[str, Any] | None:
        record = self._read(binding.task_id)
        if record is None:
            return None
        if (
            record.task_id != binding.task_id
            or record.job_id != binding.job_id
            or record.generation != binding.generation
        ):
            raise ExecutionPlanError("EXECUTION_PLAN_BINDING_CONFLICT")
        if record.source_head_oid != self._source_head(binding):
            raise ExecutionPlanError("EXECUTION_PLAN_SOURCE_HEAD_DRIFT")
        envelope = project_execution_envelope(binding, record.execution)
        return {
            "execution": dict(record.execution),
            "plan_digest": record.plan_digest,
            "execution_digest": envelope["execution_digest"],
        }