from __future__ import annotations

from dataclasses import asdict
import json
import random
import time
from typing import Any, Callable

from .documents import OfflineDocumentAdapter
from .evaluator import Evaluation, IndependentEvaluator
from .models import AgentAction, Budget, CheckpointIncompatible, ExternalStateUnknown, PolicyError, RetriableError, State, ValidationError, canonical_hash
from .policy import ALL_TOOLS, PolicyGateway
from .storage import SCHEMA_VERSION, Store
from .tools import FaultInjector, ToolExecutor


ALLOWED_TRANSITIONS: dict[State, set[State]] = {
    State.RECEIVED: {State.PARSING},
    State.PARSING: {State.EXTRACTING},
    State.EXTRACTING: {State.VALIDATING},
    State.VALIDATING: {State.NEEDS_INFORMATION, State.MANUAL_REVIEW, State.POLICY_CHECK},
    State.POLICY_CHECK: {State.READY_FOR_REVIEW, State.MANUAL_REVIEW},
    State.READY_FOR_REVIEW: {State.APPROVAL_PENDING, State.MANUAL_REVIEW},
    State.APPROVAL_PENDING: {State.APPROVED, State.MANUAL_REVIEW, State.CANCELLED},
    State.APPROVED: {State.RETRY_PENDING, State.COMPLETED, State.MANUAL_REVIEW, State.PERMANENT_FAILURE},
    State.RETRY_PENDING: {State.APPROVED, State.MANUAL_REVIEW, State.PERMANENT_FAILURE},
}


class AgentAdapter:
    def choose(self, payload: dict[str, Any]) -> AgentAction:
        return AgentAction.from_payload(payload, ALL_TOOLS)


class ScriptedAgent(AgentAdapter):
    def __init__(self, seed: int = 0) -> None:
        self.seed = seed
        self.variant = random.Random(seed).randrange(3)

    def choose(self, context: dict[str, Any]) -> AgentAction:
        case_id, run_id, phase = context["case_id"], context["run_id"], context["phase"]
        target_args = {"case_id": case_id, "target": "synthetic-review-record", "value": "REVIEW_RECORDED", "expected_version": 0}
        if phase == "draft":
            payload = {"tool": "save_review_draft", "arguments": {"case_id": case_id, "draft_version": 1 + self.variant}, "request_id": f"{run_id}-draft-v{self.variant}", "idempotency_key": f"{run_id}-draft-v{self.variant}"}
        elif phase == "approval_request":
            payload = {"tool": "request_approval", "arguments": {"case_id": case_id, "approval_id": f"approval-{run_id}", "requested_tool": "apply_mock_update", "requested_args": target_args}, "request_id": f"{run_id}-approval-request", "idempotency_key": f"{run_id}-approval-request"}
        elif phase == "apply":
            payload = {"tool": "apply_mock_update", "arguments": target_args, "request_id": f"{run_id}-apply", "idempotency_key": f"{run_id}-apply"}
        else:
            raise ValidationError("unknown scripted agent phase")
        return super().choose(payload)


class WorkflowRuntime:
    versions = {"model": "offline-v1", "prompt": "prompt-v1", "policy": "policy-v1", "schema": "schema-v1"}

    def __init__(self, store: Store | None = None, run_id: str = "run-1", faults: FaultInjector | None = None, budget: Budget | None = None, agent: AgentAdapter | None = None, clock: Callable[[], float] | None = None, sleeper: Callable[[float], None] | None = None, deadline_seconds: float = 30.0, capabilities: dict[str, bool] | None = None, seed: int = 0) -> None:
        self._owns_store = store is None
        self._closed = False
        self.store = store if store is not None else Store()
        self.run_id = run_id
        self.seed = seed
        self.clock = clock or time.monotonic
        self.sleeper = sleeper or time.sleep
        self.deadline_at = self.clock() + deadline_seconds
        self.capabilities = capabilities or {"checkpoint": True, "retry": True, "idempotency": True, "state_policy": True, "external_approval": True, "independent_evaluation": True}
        self.adapter = OfflineDocumentAdapter()
        self.policy = PolicyGateway(self.store, self.versions["policy"], self.clock)
        self.faults = faults or FaultInjector(seed)
        self.evaluator = IndependentEvaluator()
        self.agent = agent or ScriptedAgent(seed)
        self.budget = budget or Budget()
        self.tools = ToolExecutor(self.store, self.policy, run_id, self.faults, self.budget, self.clock, self.capabilities.get("idempotency", True), self.capabilities.get("state_policy", True))
        self.tools.budget = self.budget
        self.completed_steps: list[str] = []
        self.backoff_history: list[float] = []
        self.pending_approval_id: str | None = None
        self.input_hash: str | None = None
        self.active_case_id: str | None = None
        self.retry_attempt_count = 0

    @property
    def owns_store(self) -> bool:
        return self._owns_store

    @property
    def closed(self) -> bool:
        return self._closed

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        if self._owns_store:
            self.store.close()

    def __enter__(self) -> "WorkflowRuntime":
        self._ensure_open()
        return self

    def __exit__(self, exc_type, exc, traceback) -> None:
        self.close()

    def __del__(self) -> None:
        try:
            self.close()
        except Exception:
            pass

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("workflow runtime is closed")

    def transition(self, case_id: str, target: State, reason: str, approval_id: str | None = None, completion_token: str | None = None, evaluation_digest: str | None = None) -> None:
        self._ensure_open()
        current = State(self.store.state(case_id, self.run_id))
        if target not in ALLOWED_TRANSITIONS.get(current, set()):
            raise ValidationError(f"illegal transition {current.value}->{target.value}")
        if target == State.APPROVED:
            try:
                self.policy.validate_approval(approval_id, case_id, self.run_id, "apply_mock_update", self._target_args(case_id), allow_used=current == State.RETRY_PENDING)
            except PolicyError as error:
                raise ValidationError("approved transition requires current matching external approval") from error
        if self.clock() > self.deadline_at:
            self.store.safety_block(self.run_id, case_id, "DEADLINE_PREFLIGHT", {"target": target.value}, self.clock())
            raise PolicyError("deadline exceeded before transition")
        try:
            self.budget.preflight_step()
        except Exception as error:
            self.store.safety_block(self.run_id, case_id, "STEP_PREFLIGHT", {"target": target.value, "error_class": type(error).__name__}, self.clock())
            raise
        created_at = self.clock()
        if target == State.COMPLETED:
            if not completion_token or not evaluation_digest:
                raise ValidationError("completion requires case/run evaluation token")
            self.store.complete_with_token(
                completion_token,
                self.run_id,
                case_id,
                evaluation_digest,
                lambda: self.store.set_state(case_id, current.value, target.value, reason, self.run_id, created_at),
            )
        else:
            self.store.set_state(case_id, current.value, target.value, reason, self.run_id, created_at)
        self.completed_steps.append(target.value)

    def checkpoint(self, case_id: str) -> dict[str, Any]:
        self._ensure_open()
        journal = self.store.rows_for_case("journal", case_id, self.run_id)
        payload = {
            "run_id": self.run_id,
            "case_id": case_id,
            "state": self.store.state(case_id, self.run_id),
            "completed_steps": list(self.completed_steps),
            "tool_call_ids": [row["request_id"] for row in journal if row["request_id"]],
            "external_effects": [row["idem_key"] for row in self.store.rows_for_case("effects", case_id, self.run_id)],
            "versions": dict(self.versions),
            "store_schema_version": SCHEMA_VERSION,
            "budget": self.budget.to_dict(),
            "deadline_at": self.deadline_at,
            "pending_approval_id": self.pending_approval_id,
            "input_hash": self.input_hash,
            "journal_position": max((row["id"] for row in journal), default=0),
            "backoff_history": list(self.backoff_history),
        }
        self.store.save_checkpoint(self.run_id, case_id, payload)
        return payload

    def resume(self, checkpoint: dict[str, Any], documents: list[dict[str, Any]], expected_case_id: str | None = None) -> None:
        self._ensure_open()
        expected_case = expected_case_id or self.active_case_id
        if checkpoint.get("run_id") != self.run_id:
            raise CheckpointIncompatible("checkpoint run ownership mismatch")
        if not expected_case or checkpoint.get("case_id") != expected_case:
            raise CheckpointIncompatible("checkpoint case ownership mismatch")
        persisted = self.store.checkpoint(self.run_id, expected_case)
        if persisted is None or canonical_hash(persisted) != canonical_hash(checkpoint):
            raise CheckpointIncompatible("checkpoint key and payload ownership mismatch")
        if checkpoint.get("versions") != self.versions:
            raise CheckpointIncompatible("checkpoint versions are incompatible")
        if checkpoint.get("store_schema_version") != SCHEMA_VERSION:
            raise CheckpointIncompatible("store schema version is incompatible")
        if checkpoint.get("input_hash") != canonical_hash(documents):
            raise CheckpointIncompatible("checkpoint input hash changed")
        if checkpoint.get("state") != self.store.state(expected_case, self.run_id):
            raise CheckpointIncompatible("checkpoint state differs from persisted state")
        self.active_case_id = expected_case
        self.completed_steps = list(checkpoint["completed_steps"])
        self.budget = Budget.from_dict(checkpoint["budget"])
        self.tools.budget = self.budget
        self.deadline_at = float(checkpoint["deadline_at"])
        self.pending_approval_id = checkpoint.get("pending_approval_id")
        self.input_hash = checkpoint["input_hash"]
        self.backoff_history = list(checkpoint.get("backoff_history", []))

    def _agent_action(self, case_id: str, phase: str) -> AgentAction:
        self._ensure_open()
        try:
            self.budget.preflight_agent_call()
        except Exception as error:
            self.store.safety_block(self.run_id, case_id, "AGENT_PREFLIGHT", {"phase": phase, "error_class": type(error).__name__}, self.clock())
            raise
        try:
            action = self.agent.choose({"case_id": case_id, "run_id": self.run_id, "phase": phase, "state": self.store.state(case_id, self.run_id)})
        except Exception:
            self.budget.record_agent_signature(canonical_hash({"exception_phase": phase}))
            raise
        if not isinstance(action, AgentAction):
            self.budget.record_agent_signature(canonical_hash({"invalid_phase": phase, "type": type(action).__name__}))
            raise ValidationError("agent returned non-structured action")
        signature = canonical_hash({"tool": action.tool, "arguments": action.arguments})
        self.budget.record_agent_signature(signature)
        self.store.log(self.run_id, case_id, "AGENT_ACTION", {"signature": signature, "validated": True, "phase": phase}, action.request_id, action.tool, self.clock())
        return action

    def _execute_agent_action(self, case_id: str, action: AgentAction, approval_id: str | None = None) -> dict[str, Any]:
        return self.tools.execute(case_id, case_id, action.tool, action.arguments, action.request_id, action.idempotency_key, approval_id)

    def _execute_safe_draft(self, case_id: str, fault_retry_limit: int) -> bool:
        action = self._agent_action(case_id, "draft")
        attempts = 0
        while True:
            try:
                self._execute_agent_action(case_id, action)
                return True
            except (RetriableError, ExternalStateUnknown) as error:
                if not self.capabilities.get("retry", True) or attempts >= fault_retry_limit:
                    self.store.log(self.run_id, case_id, "DRAFT_RETRY_DISABLED", {"failure_class":error.category.value,"safe_preapproval_write":True}, action.request_id, action.tool, self.clock())
                    return False
                attempts += 1
                self.retry_attempt_count += 1
                backoff = 0.001 * (2 ** (attempts - 1))
                self.backoff_history.append(backoff)
                self.store.log(self.run_id, case_id, "DRAFT_RETRY_PENDING", {"failure_class":error.category.value,"retry_attempt":attempts,"backoff":backoff,"safe_preapproval_write":True}, action.request_id, action.tool, self.clock())
                if self.capabilities.get("checkpoint", False):
                    self.checkpoint(case_id)
                self.sleeper(backoff)

    def execute(self, case_id: str, documents: list[dict[str, Any]], fault_retry_limit: int = 2, stop_after_ready: bool = False) -> dict[str, Any]:
        self._ensure_open()
        self.active_case_id = case_id
        self.input_hash = canonical_hash(documents)
        preflight = {
            "deadline_expired": self.clock() > self.deadline_at,
            "max_steps": self.budget.max_steps,
            "max_calls": self.budget.max_calls,
            "max_agent_calls": self.budget.max_agent_calls,
            "max_tool_calls": self.budget.max_tool_calls,
            "max_cost_units": self.budget.max_cost_units,
        }
        if preflight["deadline_expired"] or any(preflight[name] < 1 for name in ("max_steps", "max_calls", "max_agent_calls", "max_tool_calls", "max_cost_units")):
            self.store.record_run_rejection(self.run_id, case_id, "RUN_PREFLIGHT", preflight, self.clock())
            raise PolicyError("run resource preflight blocked before state change")
        self.store.create_case(case_id, State.RECEIVED.value, run_id=self.run_id)
        self.store.init_business(case_id, self.run_id, "synthetic-review-record")
        self.transition(case_id, State.PARSING, "explicit parse start")
        self.transition(case_id, State.EXTRACTING, "adapter selected")
        doc_result = self.adapter.adapt(documents)
        self.transition(case_id, State.VALIDATING, "structured extraction completed")
        if doc_result.outcome == "NEEDS_INFORMATION":
            self.transition(case_id, State.NEEDS_INFORMATION, "required evidence missing")
            return self._result(case_id, doc_result)
        if doc_result.outcome == "MANUAL_REVIEW":
            self.transition(case_id, State.MANUAL_REVIEW, "document uncertainty")
            return self._result(case_id, doc_result)
        self.transition(case_id, State.POLICY_CHECK, "document validation passed")
        self.transition(case_id, State.READY_FOR_REVIEW, "external policy passed")
        if stop_after_ready:
            self.checkpoint(case_id)
            return self._result(case_id, doc_result)
        try:
            self._execute_safe_draft(case_id, fault_retry_limit)
            if not self.capabilities.get("external_approval", True):
                if self.capabilities.get("checkpoint", False):
                    self.checkpoint(case_id)
                return self._result(case_id, doc_result)
            request_action = self._agent_action(case_id, "approval_request")
            request_result = self._execute_agent_action(case_id, request_action)
        except Exception as error:
            self.store.add_decision(self.run_id, {"action": "agent_or_tool_failure", "evidence": [], "policy": self.policy.version, "verification": "failed", "next_state": State.MANUAL_REVIEW.value, "failure_class": type(error).__name__}, case_id)
            self.transition(case_id, State.MANUAL_REVIEW, "agent or tool failure")
            raise
        self.pending_approval_id = request_result["approval_id"]
        self.transition(case_id, State.APPROVAL_PENDING, "external approval required")
        if self.capabilities.get("checkpoint", True):
            self.checkpoint(case_id)
        return self._result(case_id, doc_result)

    def _apply_after_external_approval(self, case_id: str, doc_result: Any, fault_retry_limit: int = 2) -> dict[str, Any]:
        approval_id = self.pending_approval_id
        if not approval_id:
            raise CheckpointIncompatible("pending approval identifier missing")
        self.transition(case_id, State.APPROVED, "external approval accepted", approval_id=approval_id)
        attempts = 0
        effective_retry_limit = fault_retry_limit if self.capabilities.get("retry", True) else 0
        while True:
            attempts += 1
            try:
                if self.clock() > self.deadline_at:
                    raise PolicyError("deadline exceeded")
                action = self._agent_action(case_id, "apply")
                self._execute_agent_action(case_id, action, approval_id)
                break
            except (RetriableError, ExternalStateUnknown) as error:
                self.store.log(self.run_id, case_id, "FAILED", {"failure_class": error.category.value}, f"{self.run_id}-apply", "apply_mock_update")
                if attempts > effective_retry_limit:
                    self.transition(case_id, State.MANUAL_REVIEW, "retry exhausted")
                    return self._result(case_id, doc_result)
                self.transition(case_id, State.RETRY_PENDING, "retryable tool failure")
                backoff = 0.001 * (2 ** (attempts - 1))
                self.backoff_history.append(backoff)
                self.checkpoint(case_id)
                self.sleeper(backoff)
                if self.clock() > self.deadline_at:
                    self.transition(case_id, State.MANUAL_REVIEW, "deadline exceeded during retry")
                    return self._result(case_id, doc_result)
                self.transition(case_id, State.APPROVED, "retry window reached", approval_id=approval_id)
            except PolicyError as error:
                self.store.add_decision(self.run_id, {"action": "limit_stop", "evidence": [], "policy": self.policy.version, "verification": "blocked", "next_state": State.MANUAL_REVIEW.value, "failure_class": type(error).__name__}, case_id)
                self.transition(case_id, State.MANUAL_REVIEW, "resource or repeat limit")
                return self._result(case_id, doc_result)
        required = {State.PARSING.value, State.EXTRACTING.value, State.VALIDATING.value, State.POLICY_CHECK.value, State.READY_FOR_REVIEW.value, State.APPROVED.value}
        if self.capabilities.get("independent_evaluation", True):
            evaluation = self.evaluator.evaluate_precompletion(self.store, case_id, self.run_id, required)
        else:
            evaluation = Evaluation(False, ("independent_evaluation_disabled",), self.store.duplicate_effect_count(case_id, self.run_id), 0, 0, canonical_hash({"case_id": case_id, "run_id": self.run_id, "disabled": True}))
        if evaluation.passed:
            token = canonical_hash({"case_id": case_id, "run_id": self.run_id, "evaluation": evaluation.evaluation_digest, "nonce": "single-use"})
            self.transition(case_id, State.COMPLETED, "independent evaluation passed", completion_token=token, evaluation_digest=evaluation.evaluation_digest)
        else:
            self.transition(case_id, State.PERMANENT_FAILURE, "independent evaluation failed")
        self.store.add_decision(self.run_id, {
            "action": "complete_or_fail",
            "evidence": sorted(required),
            "policy": self.policy.version,
            "verification": "passed" if evaluation.passed else "failed",
            "next_state": self.store.state(case_id, self.run_id),
            "failure_class": None if evaluation.passed else "PERMANENT",
        }, case_id)
        return self._result(case_id, doc_result)

    def continue_ready(self, case_id: str, documents: list[dict[str, Any]], fault_retry_limit: int = 2) -> dict[str, Any]:
        self._ensure_open()
        if not self.capabilities.get("checkpoint", True):
            raise CheckpointIncompatible("checkpoint capability disabled")
        checkpoint = self.store.checkpoint(self.run_id, case_id)
        if not checkpoint:
            raise CheckpointIncompatible("checkpoint missing")
        self.resume(checkpoint, documents, case_id)
        return self._apply_after_external_approval(case_id, self.adapter.adapt(documents), fault_retry_limit)

    def handle_event(self, event_id: str, result: dict[str, Any], case_id: str) -> bool:
        self._ensure_open()
        return self.store.claim_event(event_id, result, self.run_id, case_id)

    @staticmethod
    def _target_args(case_id: str) -> dict[str, Any]:
        return {"case_id": case_id, "target": "synthetic-review-record", "value": "REVIEW_RECORDED", "expected_version": 0}

    def _result(self, case_id: str, document_result: Any) -> dict[str, Any]:
        agent_signatures = [json.loads(row["detail"])["signature"] for row in self.store.rows_for_case("journal", case_id, self.run_id) if row["phase"] == "AGENT_ACTION"]
        journal = self.store.rows_for_case("journal", case_id, self.run_id)
        return {
            "run_id": self.run_id,
            "case_id": case_id,
            "state": self.store.state(case_id, self.run_id),
            "problems": document_result.problems,
            "evidence": {key: asdict(value) for key, value in document_result.fields.items()},
            "untrusted_instructions": document_result.untrusted_instructions,
            "effects": self.store.effect_count(case_id, self.run_id),
            "duplicate_effects": self.store.duplicate_effect_count(case_id, self.run_id),
            "steps": self.budget.steps,
            "calls": self.budget.calls,
            "cost_units": self.budget.cost_units,
            "agent_call_count": self.budget.agent_calls,
            "tool_call_count": self.budget.tool_calls,
            "retry_attempt_count": self.retry_attempt_count,
            "retry_pending_count": sum(row["phase"] in {"DRAFT_RETRY_PENDING"} for row in journal) + sum(row["to_state"] == State.RETRY_PENDING.value for row in self.store.rows_for_case("transitions", case_id, self.run_id)),
            "committed_result_reuse_count": sum(row["phase"] == "RESULT_REUSED" for row in journal),
            "backoff_history": list(self.backoff_history),
            "backoff_total": sum(self.backoff_history),
            "deadline_at": self.deadline_at,
            "pending_approval_id": self.pending_approval_id,
            "approval_status_counts": self._approval_counts(case_id),
            "business_state": self.store.business(case_id, self.run_id, "synthetic-review-record"),
            "capability_branches": {
                "checkpoint": self.capabilities.get("checkpoint", True), "retry": self.capabilities.get("retry", True),
                "idempotency": self.capabilities.get("idempotency", True), "state_policy": self.capabilities.get("state_policy", True),
                "external_approval": self.capabilities.get("external_approval", True), "independent_evaluation": self.capabilities.get("independent_evaluation", True),
            },
            "fault_plan_digest": self.faults.plan_digest(),
            "fault_trace": list(self.faults.trace),
            "agent_trace_digest": canonical_hash(agent_signatures),
            "agent_trace": agent_signatures,
            "seed": self.seed,
            "seed_effective": None,
        }

    def _approval_counts(self, case_id: str) -> dict[str, int]:
        counts: dict[str, int] = {}
        for row in self.store.rows_for_case("approvals", case_id, self.run_id):
            counts[row["status"]] = counts.get(row["status"], 0) + 1
        return counts


def classify_error(error: Exception) -> str:
    if hasattr(error, "category"):
        return str(error.category.value)
    return "PERMANENT"
