import asyncio
import uuid
from datetime import UTC, datetime
from typing import Annotated

from fastapi import APIRouter, Cookie, Depends, Header, Request
from pydantic import BaseModel, Field
from sqlalchemy import select

from .ai_client import AIPlanResult, PlanClient, get_ai_plan_client
from .auth import AuthDep, SessionDep, _audit, _require_csrf
from .authorization import authorize, permissions_for
from .config import get_settings
from .errors import AppError
from .health import AIHealthDep, RedisDep
from .models import (
    Agent,
    AgentPermission,
    AgentTool,
    AIModelProfile,
    DeviceSession,
    Execution,
    ExecutionStep,
)
from .policy import ActionRequest, AutonomyLevel, Decision, evaluate_action
from .security import canonical_payload_hash, privacy_hash
from .tool_runtime import CORE_TOOLS, ToolContext, execute_core_tool, get_core_tool

router = APIRouter(prefix="/api/v1/executions", tags=["executions"])
AIClientDep = Annotated[PlanClient, Depends(get_ai_plan_client)]


class CreateExecutionInput(BaseModel):
    instruction: str = Field(min_length=1, max_length=20_000)
    agent_id: uuid.UUID | None = None


def _execution_data(record: Execution, steps: list[ExecutionStep] | None = None) -> dict:
    data = {
        "id": str(record.id),
        "status": record.status,
        "agent_id": str(record.agent_id) if record.agent_id else None,
        "provider": record.provider,
        "model": record.model,
        "plan": record.plan,
        "output_summary": record.output_summary,
        "error_code": record.error_code,
        "input_tokens": record.input_tokens,
        "output_tokens": record.output_tokens,
        "stop_requested_at": record.stop_requested_at,
        "created_at": record.created_at,
        "updated_at": record.updated_at,
    }
    if steps is not None:
        data["steps"] = []
        for step in steps:
            checkpoint = step.checkpoint or {}
            data["steps"].append({
                "id": str(step.id),
                "position": step.position,
                "tool_name": step.tool_name,
                "permission": step.permission,
                "risk": step.risk,
                "access_type": step.access_type,
                "status": step.status,
                "concise_rationale": step.concise_rationale,
                "depends_on": step.depends_on,
                "result": checkpoint.get("result"),
                "verification": checkpoint.get("verification"),
                "error_code": checkpoint.get("error_code"),
            })
    return data


async def _mark_planning_error(session: SessionDep, *, execution_id: uuid.UUID,
                               workspace_id: uuid.UUID, actor_id: uuid.UUID,
                               request_id: str, error: AppError) -> None:
    record = await session.scalar(
        select(Execution).where(
            Execution.id == execution_id,
            Execution.workspace_id == workspace_id,
        ).with_for_update()
    )
    if record and record.status != "cancelled":
        record.status = "blocked" if error.status_code == 503 else "failed"
        record.error_code = error.code
        record.completed_at = datetime.now(UTC)
        record.version += 1
        await _audit(
            session,
            workspace_id=workspace_id,
            actor_id=actor_id,
            event_type="execution.planning_failed",
            request_id=request_id,
            data={"execution_id": str(execution_id), "error_code": error.code},
        )
        await session.commit()


def _validate_plan_policy(result: AIPlanResult, granted: set[str]) -> None:
    for step in result.plan.steps:
        policy = get_core_tool(step.tool_name)
        policy.validate_arguments(step.arguments)
        decision = evaluate_action(
            ActionRequest(
                permission=policy.permission,
                risk=policy.risk,
                is_write=policy.is_write,
            ),
            AutonomyLevel.OBSERVE,
            granted,
        )
        if decision != Decision.ALLOW:
            raise AppError("AI_PLAN_NOT_AUTHORIZED", "The plan was not authorized.", 403)


@router.post("", status_code=201)
async def create_execution(
    payload: CreateExecutionInput,
    request: Request,
    auth: AuthDep,
    session: SessionDep,
    ai_client: AIClientDep,
    hayva_csrf: Annotated[str | None, Cookie()] = None,
    x_csrf_token: Annotated[str | None, Header()] = None,
):
    await authorize(
        session,
        workspace_id=auth.workspace_id,
        user_id=auth.user_id,
        permission="executions.create",
    )
    device_session = await session.get(DeviceSession, auth.session_id)
    _require_csrf(device_session, hayva_csrf, x_csrf_token)
    granted = await permissions_for(
        session, workspace_id=auth.workspace_id, user_id=auth.user_id
    )
    agent_statement = select(Agent).where(
        Agent.workspace_id == auth.workspace_id,
        Agent.status == "active",
    )
    if payload.agent_id:
        agent_statement = agent_statement.where(Agent.id == payload.agent_id)
    else:
        agent_statement = agent_statement.where(
            Agent.name == "Personal Executive Assistant"
        )
    agent = await session.scalar(agent_statement)
    if not agent:
        raise AppError("AGENT_UNAVAILABLE", "An active execution agent is required.", 409)
    agent_permissions = set((await session.scalars(select(
        AgentPermission.permission_key
    ).where(
        AgentPermission.workspace_id == auth.workspace_id,
        AgentPermission.agent_id == agent.id,
    ))).all())
    agent_tools = set((await session.scalars(select(AgentTool.tool_name).where(
        AgentTool.workspace_id == auth.workspace_id,
        AgentTool.agent_id == agent.id,
    ))).all())
    allowed_tools = sorted(
        name for name, policy in CORE_TOOLS.items()
        if name in agent_tools
        and policy.permission in agent_permissions
        and policy.permission in granted
    )
    selected_model = None
    if agent.preferred_model:
        model_profile = await session.scalar(select(AIModelProfile).where(
            AIModelProfile.workspace_id == auth.workspace_id,
            AIModelProfile.profile_key == agent.preferred_model,
            AIModelProfile.enabled.is_(True),
        ))
        if not model_profile:
            raise AppError("AGENT_MODEL_UNAVAILABLE", "The agent model is unavailable.", 409)
        selected_model = model_profile.model_id
    execution = Execution(
        id=uuid.uuid4(),
        workspace_id=auth.workspace_id,
        requested_by_user_id=auth.user_id,
        agent_id=agent.id,
        instruction=payload.instruction,
        instruction_hash=canonical_payload_hash({"instruction": payload.instruction}),
        status="planning",
    )
    session.add(execution)
    await _audit(
        session,
        workspace_id=auth.workspace_id,
        actor_id=auth.user_id,
        event_type="execution.created",
        request_id=request.state.request_id,
        data={"execution_id": str(execution.id), "allowed_tools": allowed_tools},
    )
    await session.commit()

    safety_identifier = privacy_hash(
        get_settings().require_app_secret(),
        f"{auth.workspace_id}:{auth.user_id}",
    ) or ""
    try:
        result = await ai_client.create_plan(
            execution_id=execution.id,
            workspace_id=auth.workspace_id,
            actor_id=auth.user_id,
            instruction=payload.instruction,
            allowed_tools=allowed_tools,
            safety_identifier=safety_identifier,
            request_id=request.state.request_id,
            model=selected_model,
        )
        _validate_plan_policy(result, granted)
    except AppError as error:
        await _mark_planning_error(
            session,
            execution_id=execution.id,
            workspace_id=auth.workspace_id,
            actor_id=auth.user_id,
            request_id=request.state.request_id,
            error=error,
        )
        raise

    record = await session.scalar(
        select(Execution).where(
            Execution.id == execution.id,
            Execution.workspace_id == auth.workspace_id,
        ).with_for_update()
    )
    if not record:
        raise AppError("EXECUTION_NOT_FOUND", "The execution is unavailable.", 404)
    if record.status == "cancelled":
        await session.commit()
        await session.refresh(record)
        return {"success": True, "execution": _execution_data(record, [])}

    record.status = "planned"
    record.provider = result.provider
    record.model = result.model
    record.provider_response_id = result.provider_response_id
    record.plan = result.plan.model_dump(mode="json")
    record.input_tokens = result.usage.input_tokens
    record.output_tokens = result.usage.output_tokens
    record.version += 1
    steps: list[ExecutionStep] = []
    for position, proposed in enumerate(result.plan.steps):
        policy = get_core_tool(proposed.tool_name)
        step = ExecutionStep(
            id=uuid.uuid4(),
            workspace_id=auth.workspace_id,
            execution_id=record.id,
            position=position,
            model_step_id=proposed.step_id,
            tool_name=proposed.tool_name,
            permission=policy.permission,
            risk=policy.risk,
            access_type="write" if policy.is_write else "read",
            arguments=proposed.arguments,
            concise_rationale=proposed.concise_rationale,
            depends_on=proposed.depends_on,
            status="pending",
        )
        steps.append(step)
        session.add(step)
    await _audit(
        session,
        workspace_id=auth.workspace_id,
        actor_id=auth.user_id,
        event_type="execution.planned",
        request_id=request.state.request_id,
        data={"execution_id": str(record.id), "provider": result.provider,
              "model": result.model, "step_count": len(steps)},
    )
    await session.commit()
    await session.refresh(record)
    return {"success": True, "execution": _execution_data(record, steps)}


async def _execution_steps(
    session: SessionDep, *, workspace_id: uuid.UUID, execution_id: uuid.UUID,
    lock: bool = False,
) -> list[ExecutionStep]:
    statement = (
        select(ExecutionStep)
        .where(
            ExecutionStep.workspace_id == workspace_id,
            ExecutionStep.execution_id == execution_id,
        )
        .order_by(ExecutionStep.position)
    )
    if lock:
        statement = statement.with_for_update()
    return list((await session.scalars(statement)).all())


async def _mark_execution_error(
    session: SessionDep,
    *,
    execution_id: uuid.UUID,
    workspace_id: uuid.UUID,
    actor_id: uuid.UUID,
    request_id: str,
    step_id: uuid.UUID | None,
    error: AppError,
) -> None:
    record = await session.scalar(
        select(Execution).where(
            Execution.id == execution_id,
            Execution.workspace_id == workspace_id,
        ).with_for_update()
    )
    if not record or record.status == "cancelled":
        await session.commit()
        return
    if step_id:
        step = await session.scalar(
            select(ExecutionStep).where(
                ExecutionStep.id == step_id,
                ExecutionStep.workspace_id == workspace_id,
                ExecutionStep.execution_id == execution_id,
            ).with_for_update()
        )
        if step and step.status not in {"succeeded", "cancelled"}:
            checkpoint = dict(step.checkpoint or {})
            checkpoint.update({
                "error_code": error.code,
                "failed_at": datetime.now(UTC).isoformat(),
            })
            step.status = "failed"
            step.checkpoint = checkpoint
    record.status = "failed"
    record.error_code = error.code
    record.completed_at = datetime.now(UTC)
    record.version += 1
    await _audit(
        session,
        workspace_id=workspace_id,
        actor_id=actor_id,
        event_type="execution.failed",
        request_id=request_id,
        data={
            "execution_id": str(execution_id),
            "step_id": str(step_id) if step_id else None,
            "error_code": error.code,
        },
    )
    await session.commit()


def _authorize_runtime_step(
    step: ExecutionStep,
    granted: set[str],
    agent: Agent,
    agent_tools: set[str],
    agent_permissions: set[str],
):
    tool = get_core_tool(step.tool_name)
    if (
        step.permission != tool.permission
        or step.risk != tool.risk
        or step.access_type != ("write" if tool.is_write else "read")
    ):
        raise AppError("TOOL_POLICY_MISMATCH", "Trusted tool policy did not match the plan.", 409)
    if agent.status != "active" or agent.autonomy_level == 0:
        raise AppError("AGENT_DISABLED", "The execution agent is disabled.", 409)
    if step.tool_name not in agent_tools or tool.permission not in agent_permissions:
        raise AppError("AGENT_TOOL_REVOKED", "The agent tool grant is no longer active.", 403)
    tool.validate_arguments(step.arguments)
    decision = evaluate_action(
        ActionRequest(permission=tool.permission, risk=tool.risk, is_write=tool.is_write),
        AutonomyLevel(agent.autonomy_level),
        granted,
    )
    if decision != Decision.ALLOW:
        raise AppError("AI_PLAN_NOT_AUTHORIZED", "The plan is no longer authorized.", 403)
    return tool


@router.post("/{execution_id}/run")
async def run_execution(
    execution_id: uuid.UUID,
    request: Request,
    auth: AuthDep,
    session: SessionDep,
    redis: RedisDep,
    ai_probe: AIHealthDep,
    hayva_csrf: Annotated[str | None, Cookie()] = None,
    x_csrf_token: Annotated[str | None, Header()] = None,
):
    await authorize(
        session,
        workspace_id=auth.workspace_id,
        user_id=auth.user_id,
        permission="executions.run",
    )
    device_session = await session.get(DeviceSession, auth.session_id)
    _require_csrf(device_session, hayva_csrf, x_csrf_token)
    record = await session.scalar(
        select(Execution).where(
            Execution.id == execution_id,
            Execution.workspace_id == auth.workspace_id,
        ).with_for_update()
    )
    if not record:
        raise AppError("EXECUTION_NOT_FOUND", "The execution is unavailable.", 404)
    if record.requested_by_user_id != auth.user_id:
        raise AppError("EXECUTION_ACTOR_MISMATCH", "The execution belongs to another actor.", 403)
    if record.status == "succeeded":
        steps = await _execution_steps(
            session, workspace_id=auth.workspace_id, execution_id=execution_id
        )
        await session.commit()
        return {"success": True, "execution": _execution_data(record, steps)}
    if record.status != "planned":
        raise AppError("EXECUTION_NOT_RUNNABLE", "The execution cannot be run in its state.", 409)
    record.status = "running"
    record.started_at = datetime.now(UTC)
    record.error_code = None
    record.version += 1
    await _audit(
        session,
        workspace_id=auth.workspace_id,
        actor_id=auth.user_id,
        event_type="execution.started",
        request_id=request.state.request_id,
        data={"execution_id": str(record.id)},
    )
    await session.commit()

    while True:
        record = await session.scalar(
            select(Execution).where(
                Execution.id == execution_id,
                Execution.workspace_id == auth.workspace_id,
            ).with_for_update()
        )
        if not record:
            raise AppError("EXECUTION_NOT_FOUND", "The execution is unavailable.", 404)
        steps = await _execution_steps(
            session, workspace_id=auth.workspace_id, execution_id=execution_id, lock=True
        )
        if record.status == "cancelled":
            await session.commit()
            return {"success": True, "execution": _execution_data(record, steps)}
        succeeded = {step.model_step_id for step in steps if step.status == "succeeded"}
        pending = [step for step in steps if step.status == "pending"]
        if not pending:
            record.status = "succeeded"
            record.completed_at = datetime.now(UTC)
            record.output_summary = f"Completed {len(steps)} verified read step(s)."
            record.version += 1
            await _audit(
                session,
                workspace_id=auth.workspace_id,
                actor_id=auth.user_id,
                event_type="execution.succeeded",
                request_id=request.state.request_id,
                data={"execution_id": str(record.id), "step_count": len(steps)},
            )
            await session.commit()
            await session.refresh(record)
            return {"success": True, "execution": _execution_data(record, steps)}
        step = next(
            (candidate for candidate in pending if set(candidate.depends_on) <= succeeded),
            None,
        )
        if not step:
            error = AppError("EXECUTION_GRAPH_BLOCKED", "The execution graph cannot progress.", 409)
            await session.rollback()
            await _mark_execution_error(
                session,
                execution_id=execution_id,
                workspace_id=auth.workspace_id,
                actor_id=auth.user_id,
                request_id=request.state.request_id,
                step_id=None,
                error=error,
            )
            raise error

        runtime_step_id = step.id
        try:
            granted = await permissions_for(
                session, workspace_id=auth.workspace_id, user_id=auth.user_id
            )
            agent = await session.scalar(select(Agent).where(
                Agent.id == record.agent_id,
                Agent.workspace_id == auth.workspace_id,
            ))
            if not agent:
                raise AppError("AGENT_UNAVAILABLE", "The execution agent is unavailable.", 409)
            agent_tools = set((await session.scalars(select(AgentTool.tool_name).where(
                AgentTool.workspace_id == auth.workspace_id,
                AgentTool.agent_id == agent.id,
            ))).all())
            agent_permissions = set((await session.scalars(select(
                AgentPermission.permission_key
            ).where(
                AgentPermission.workspace_id == auth.workspace_id,
                AgentPermission.agent_id == agent.id,
            ))).all())
            tool = _authorize_runtime_step(
                step, granted, agent, agent_tools, agent_permissions
            )
        except AppError as error:
            await session.rollback()
            await _mark_execution_error(
                session,
                execution_id=execution_id,
                workspace_id=auth.workspace_id,
                actor_id=auth.user_id,
                request_id=request.state.request_id,
                step_id=runtime_step_id,
                error=error,
            )
            raise

        started_at = datetime.now(UTC)
        step.status = "running"
        step.checkpoint = {"attempt": 1, "started_at": started_at.isoformat()}
        await _audit(
            session,
            workspace_id=auth.workspace_id,
            actor_id=auth.user_id,
            event_type="execution.step_started",
            request_id=request.state.request_id,
            data={
                "execution_id": str(record.id),
                "step_id": str(step.id),
                "tool_name": tool.name,
            },
        )
        await session.commit()

        context = ToolContext(
            session=session,
            redis=redis,
            ai_probe=ai_probe,
            workspace_id=auth.workspace_id,
            user_id=auth.user_id,
        )
        try:
            async with asyncio.timeout(tool.timeout_seconds):
                result = await execute_core_tool(tool, step.arguments, context)
        except TimeoutError:
            error = AppError("TOOL_TIMEOUT", "The tool timed out.", 504)
        except AppError as tool_error:
            error = tool_error
        except Exception:  # noqa: BLE001 - execution boundary must persist a safe failure
            error = AppError("TOOL_EXECUTION_FAILED", "The tool failed safely.", 500)
        else:
            error = None
        if error:
            await session.rollback()
            await _mark_execution_error(
                session,
                execution_id=execution_id,
                workspace_id=auth.workspace_id,
                actor_id=auth.user_id,
                request_id=request.state.request_id,
                step_id=runtime_step_id,
                error=error,
            )
            raise error

        record = await session.scalar(
            select(Execution).where(
                Execution.id == execution_id,
                Execution.workspace_id == auth.workspace_id,
            ).with_for_update()
        )
        current_step = await session.scalar(
            select(ExecutionStep).where(
                ExecutionStep.id == runtime_step_id,
                ExecutionStep.workspace_id == auth.workspace_id,
                ExecutionStep.execution_id == execution_id,
            ).with_for_update()
        )
        if not record or not current_step:
            raise AppError("EXECUTION_NOT_FOUND", "The execution is unavailable.", 404)
        if record.status == "cancelled":
            current_step.status = "cancelled"
            current_step.checkpoint = {
                "attempt": 1,
                "cancelled_at": datetime.now(UTC).isoformat(),
            }
            await session.commit()
            steps = await _execution_steps(
                session, workspace_id=auth.workspace_id, execution_id=execution_id
            )
            return {"success": True, "execution": _execution_data(record, steps)}
        verified_at = datetime.now(UTC)
        current_step.status = "succeeded"
        current_step.checkpoint = {
            "attempt": 1,
            "result": result,
            "verification": {
                "status": "verified",
                "strategy": tool.verification_strategy,
                "verified_at": verified_at.isoformat(),
            },
        }
        await _audit(
            session,
            workspace_id=auth.workspace_id,
            actor_id=auth.user_id,
            event_type="execution.step_succeeded",
            request_id=request.state.request_id,
            data={
                "execution_id": str(record.id),
                "step_id": str(current_step.id),
                "tool_name": tool.name,
                "verification_strategy": tool.verification_strategy,
            },
        )
        await session.commit()


@router.get("")
async def list_executions(auth: AuthDep, session: SessionDep):
    await authorize(
        session,
        workspace_id=auth.workspace_id,
        user_id=auth.user_id,
        permission="executions.read",
    )
    records = (await session.scalars(
        select(Execution)
        .where(Execution.workspace_id == auth.workspace_id)
        .order_by(Execution.created_at.desc())
        .limit(100)
    )).all()
    return {"success": True, "executions": [_execution_data(record) for record in records]}


@router.get("/{execution_id}")
async def get_execution(execution_id: uuid.UUID, auth: AuthDep, session: SessionDep):
    await authorize(
        session,
        workspace_id=auth.workspace_id,
        user_id=auth.user_id,
        permission="executions.read",
    )
    record = await session.scalar(select(Execution).where(
        Execution.id == execution_id, Execution.workspace_id == auth.workspace_id))
    if not record:
        raise AppError("EXECUTION_NOT_FOUND", "The execution is unavailable.", 404)
    steps = (await session.scalars(select(ExecutionStep).where(
        ExecutionStep.workspace_id == auth.workspace_id,
        ExecutionStep.execution_id == execution_id).order_by(ExecutionStep.position))).all()
    return {"success": True, "execution": _execution_data(record, list(steps))}


@router.post("/{execution_id}/cancel")
async def cancel_execution(execution_id: uuid.UUID, request: Request, auth: AuthDep,
                           session: SessionDep,
                           hayva_csrf: Annotated[str | None, Cookie()] = None,
                           x_csrf_token: Annotated[str | None, Header()] = None):
    await authorize(
        session,
        workspace_id=auth.workspace_id,
        user_id=auth.user_id,
        permission="executions.cancel",
    )
    device_session = await session.get(DeviceSession, auth.session_id)
    _require_csrf(device_session, hayva_csrf, x_csrf_token)
    record = await session.scalar(select(Execution).where(
        Execution.id == execution_id,
        Execution.workspace_id == auth.workspace_id).with_for_update())
    if not record:
        raise AppError("EXECUTION_NOT_FOUND", "The execution is unavailable.", 404)
    if record.status in {"succeeded", "failed"}:
        raise AppError("EXECUTION_TERMINAL", "The execution has already finished.", 409)
    if record.status != "cancelled":
        now = datetime.now(UTC)
        record.status = "cancelled"
        record.stop_requested_at = now
        record.completed_at = now
        record.version += 1
        steps = await _execution_steps(
            session,
            workspace_id=auth.workspace_id,
            execution_id=execution_id,
            lock=True,
        )
        for step in steps:
            if step.status in {"pending", "running"}:
                step.status = "cancelled"
                checkpoint = dict(step.checkpoint or {})
                checkpoint["cancelled_at"] = now.isoformat()
                step.checkpoint = checkpoint
        await _audit(
            session,
            workspace_id=auth.workspace_id,
            actor_id=auth.user_id,
            event_type="execution.cancelled",
            request_id=request.state.request_id,
            data={"execution_id": str(record.id)},
        )
        await session.commit()
        await session.refresh(record)
    return {"success": True, "execution": _execution_data(record)}