"""Plan executor for Autopilot (policy-enforced)."""

from __future__ import annotations

import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from .contracts import PlanStep, PlanV1, ReceiptChange, ReceiptV1
from .policy import AutonomyPolicy, compute_policy_hash, evaluate_policy
from .tools import AutopilotToolDefinition, ToolRegistry
from .verification import run_postchecks, run_prechecks, run_verification


EXECUTOR_VERSION = "0.1.0"


class PlanExecutor:
    def __init__(self, policy: AutonomyPolicy, registry: ToolRegistry, state_dir: Path):
        self.policy = policy
        self.registry = registry
        self.state_dir = Path(state_dir)
        self._idempotency_seen: dict[str, set[str]] = {}
        self._sleep = time.sleep

    def execute(
        self,
        plan: PlanV1,
        approve_token: str | None = None,
        *,
        dry_run: bool = False,
    ) -> ReceiptV1:
        started_at = datetime.now(timezone.utc).isoformat()
        decision_path: list[str] = []
        changes: list[ReceiptChange] = []
        executed_any = False
        denied_any = False
        skipped_any = False

        plan.policy_hash = compute_policy_hash(self.policy)
        policy_hash = plan.policy_hash

        total_records = 0
        steps_attempted = 0
        outcome = "success"

        for step in plan.steps:
            tool_def = self.registry.get_tool(step.tool)
            decision = evaluate_policy(self.policy, step, tool_def)
            decision_path.extend(decision.decision_path)

            if decision.action == "deny":
                denied_any = True
                decision_path.append(f"step:{step.step_id}:denied")
                continue

            if decision.action == "approval_required":
                if not approve_token:
                    outcome = "approval_required"
                    decision_path.append(f"step:{step.step_id}:approval_required")
                    break
                decision_path.append(f"step:{step.step_id}:approval_granted")

            record_count = _estimate_record_count(step)
            if record_count > self.policy.defaults.max_bulk_batch:
                outcome = "budget_exceeded"
                decision_path.append("budget:max_bulk_batch")
                break
            if total_records + record_count > self.policy.defaults.max_records_per_run:
                outcome = "budget_exceeded"
                decision_path.append("budget:max_records_per_run")
                break

            if step.idempotency_key:
                seen = self._idempotency_seen.setdefault(plan.run_id, set())
                if step.idempotency_key in seen:
                    skipped_any = True
                    decision_path.append(f"step:{step.step_id}:idempotent_skip")
                    continue
                seen.add(step.idempotency_key)

            if steps_attempted >= self.policy.defaults.max_steps_per_run:
                outcome = "budget_exceeded"
                decision_path.append("budget:max_steps_per_run")
                break
            steps_attempted += 1

            if step.diff_preview is not None and step.args.get("simulate_live_mismatch"):
                decision_path.append("precheck:live_state_mismatch")

            prechecks = run_prechecks(step, tool_def)
            if any(not check.passed for check in prechecks):
                outcome = "precheck_failed"
                decision_path.append(f"step:{step.step_id}:precheck_failed")
                break

            try:
                if dry_run:
                    result = self._execute_step_dry_run(step, tool_def)
                else:
                    result = self._execute_with_retry(step, tool_def)
            except Exception as exc:
                partial_successes = _extract_partial_successes(exc)
                if partial_successes:
                    changes.append(
                        ReceiptChange(
                            tool=step.tool,
                            targets=_extract_partial_targets(partial_successes, step),
                            before=step.diff_preview,
                            after={"partial_successes": partial_successes},
                        )
                    )
                    executed_any = True
                    decision_path.append(
                        f"step:{step.step_id}:partial_success:{len(partial_successes)}"
                    )
                outcome = "partial_failure" if executed_any else "failed"
                decision_path.append(f"step:{step.step_id}:error:{exc}")
                break

            postchecks = run_postchecks(step, tool_def)
            if any(not check.passed for check in postchecks):
                outcome = "partial_failure"
                decision_path.append(f"step:{step.step_id}:postcheck_failed")
                break

            change = ReceiptChange(
                tool=step.tool,
                targets=result.get("targets", _extract_targets(step)),
                before=result.get("before", step.diff_preview),
                after=result.get("after", {"status": "executed"}),
            )
            changes.append(change)
            executed_any = True
            total_records += record_count

        if outcome == "success":
            if executed_any:
                outcome = "success"
            elif denied_any:
                outcome = "denied"
            elif skipped_any:
                outcome = "skipped"
            else:
                outcome = "no_op"

        verification = run_verification(plan, self.registry, policy=self.policy, state_dir=self.state_dir)

        completed_at = datetime.now(timezone.utc).isoformat()
        receipt = ReceiptV1(
            run_id=plan.run_id,
            policy_hash=policy_hash,
            plan_hash=plan.plan_hash,
            executor_version=EXECUTOR_VERSION,
            tool_versions={"g-gremlin": EXECUTOR_VERSION},
            changes=changes,
            verification=verification,
            artifacts=[],
            decision_path=decision_path,
            started_at=started_at,
            completed_at=completed_at,
            outcome=outcome,
        )
        return receipt

    def _execute_with_retry(
        self,
        step: PlanStep,
        tool_def: AutopilotToolDefinition | None,
    ) -> dict[str, Any]:
        # Max retries is the number of retries after the initial attempt.
        # With 3 retries and delays [1, 2, 4], total attempts = 4.
        max_retries = 3
        delays = [1, 2, 4]
        attempt = 0
        while True:
            attempt += 1
            try:
                return self._execute_step(step, tool_def)
            except Exception as exc:
                if _is_lock_error(exc) and attempt <= max_retries:
                    self._sleep(delays[attempt - 1])
                    continue
                raise

    def _execute_step(
        self,
        step: PlanStep,
        tool_def: AutopilotToolDefinition | None,
    ) -> dict[str, Any]:
        simulated = step.args.get("simulate_error")
        if simulated == "lock":
            raise RuntimeError("UNABLE_TO_LOCK_ROW")
        if simulated == "fail":
            raise RuntimeError("simulated failure")
        handler = self.registry.get_handler(step.tool)
        if handler is None:
            raise RuntimeError(f"No handler registered for tool: {step.tool}")
        return handler(step, tool_def)

    def _execute_step_dry_run(
        self,
        step: PlanStep,
        tool_def: AutopilotToolDefinition | None,
    ) -> dict[str, Any]:
        return {
            "targets": _extract_targets(step),
            "before": step.diff_preview,
            "after": {"status": "dry_run", "tool": step.tool},
        }


def _is_lock_error(exc: Exception) -> bool:
    msg = str(exc).lower()
    return "unable_to_lock_row" in msg or "unable to obtain exclusive access" in msg


def _extract_targets(step: PlanStep) -> list[str]:
    args = step.args or {}
    for key in ("record_ids", "user_ids", "targets", "ids"):
        if isinstance(args.get(key), list):
            return [str(x) for x in args.get(key, [])]
    if isinstance(args.get("records"), list):
        return [str(row.get("record_id")) for row in args.get("records", []) if row.get("record_id")]
    return []


def _estimate_record_count(step: PlanStep) -> int:
    args = step.args or {}
    for key in ("record_ids", "user_ids", "targets", "ids"):
        if isinstance(args.get(key), list):
            return len(args.get(key))
    if isinstance(args.get("records"), list):
        return len(args.get("records"))
    count = args.get("count")
    if count is not None:
        try:
            return int(count)
        except Exception:
            return 1
    return 1


def _extract_partial_successes(exc: Exception) -> list[dict[str, Any]]:
    successes = getattr(exc, "successes", None)
    if not successes:
        return []
    if isinstance(successes, list):
        return [item for item in successes if isinstance(item, dict)]
    return []


def _extract_partial_targets(successes: list[dict[str, Any]], step: PlanStep) -> list[str]:
    targets: list[str] = []
    for item in successes:
        user_id = item.get("user_id")
        perm_id = item.get("permission_set_id") or item.get("permission_set")
        if user_id and perm_id:
            targets.append(f"{user_id}:{perm_id}")
        elif user_id:
            targets.append(str(user_id))
    if targets:
        return targets
    return _extract_targets(step)


__all__ = ["PlanExecutor", "EXECUTOR_VERSION"]
