"""Reliable pull-mode Step lease lifecycle for Host Agents."""
from __future__ import annotations

import datetime as _dt
from typing import Any, Dict, List, Optional

from sqlalchemy import delete, func, select
from sqlalchemy.ext.asyncio import AsyncSession

from ..config import get_settings
from ..enums import (
    EventKind,
    EventSource,
    OnFail,
    RunState,
    Severity,
    StepState,
    Verdict,
)
from ..events import append_event
from ..evidence import build_evidence_bundle
from ..models import ResourceLock, Run, RunEvent, RunStep, TestHost, iso_z, utcnow
from ..schemas import AgentStepCompleteBody, AgentStepEvent

_ACTIVE_STEP_STATES = {
    StepState.PENDING.value,
    StepState.DISPATCHED.value,
    StepState.RUNNING.value,
}


class StepLeaseConflict(ValueError):
    """The Agent request does not match the current lease generation/state."""


def step_status(step: RunStep, *, replayed: bool = False) -> Dict[str, Any]:
    return {
        "step_id": step.id,
        "run_id": step.run_id,
        "state": step.state,
        "attempt": step.attempt,
        "lease_expires_at": iso_z(step.lease_expires_at),
        "replayed": replayed,
    }


def step_lease(step: RunStep, run: Run) -> Dict[str, Any]:
    return {
        "step_id": step.id,
        "run_id": step.run_id,
        "seq": step.seq,
        "loop_index": step.loop_index,
        "step_key": step.step_key,
        "step_type": step.step_type,
        "adapter": step.adapter,
        "template": step.template,
        "params": step.params,
        "danger": step.danger,
        "idempotent": step.idempotent,
        "timeout_s": step.timeout_s,
        "attempt": step.attempt,
        "lease_expires_at": iso_z(step.lease_expires_at),
        # This slice exercises the real process/network protocol but never executes
        # a real hardware command.  A future signed adapter registry will remove it.
        "dry_run": run.params.get("execution_mode") == "agent_dry_run",
    }


async def _locked_run(session: AsyncSession, run_id: str) -> Run:
    run = await session.scalar(
        select(Run).where(Run.id == run_id).with_for_update()
    )
    if run is None:
        raise StepLeaseConflict("Run 不存在")
    return run


async def _owned_step(
    session: AsyncSession, host: TestHost, step_id: str
) -> tuple[RunStep, Run]:
    step = await session.scalar(
        select(RunStep)
        .join(Run, Run.id == RunStep.run_id)
        .where(RunStep.id == step_id, Run.test_host_id == host.id)
        .with_for_update()
    )
    if step is None:
        raise StepLeaseConflict("步骤不存在或不属于当前 Agent")
    run = await _locked_run(session, step.run_id)
    return step, run


async def _replayed(
    session: AsyncSession,
    *,
    run_id: str,
    step_id: str,
    idempotency_key: str,
    operation: str,
) -> bool:
    event = await session.scalar(
        select(RunEvent).where(
            RunEvent.run_id == run_id,
            RunEvent.idempotency_key == idempotency_key,
        )
    )
    if event is None:
        return False
    if event.step_id != step_id or event.payload.get("_operation") != operation:
        raise StepLeaseConflict("Idempotency-Key 已用于其他步骤或操作")
    return True


def _assert_current(step: RunStep, agent_id: str, attempt: int) -> None:
    now = utcnow()
    if step.leased_by != agent_id or step.attempt != attempt:
        raise StepLeaseConflict("租约代次不匹配，步骤可能已被重新派发")
    if step.lease_expires_at is None or step.lease_expires_at <= now:
        raise StepLeaseConflict("步骤租约已过期")
    if (
        step.state == StepState.RUNNING.value
        and step.started_at is not None
        and step.started_at + _dt.timedelta(seconds=step.timeout_s) <= now
    ):
        raise StepLeaseConflict("步骤执行已超过 timeout_s")


async def _freeze_run(
    session: AsyncSession, run: Run, step: RunStep, reason: str
) -> None:
    if run.state not in {RunState.RUNNING.value, RunState.RECOVERING.value}:
        return
    await append_event(
        session,
        run,
        EventKind.RUN_FROZEN.value,
        reason,
        severity=Severity.CRITICAL.value,
        target_state=RunState.FROZEN.value,
        projection={
            "verdict": Verdict.INCONCLUSIVE.value,
            "freeze_reason": reason,
            "ended_at": utcnow(),
        },
        step_id=step.id,
        loop_index=step.loop_index,
    )
    await session.execute(delete(ResourceLock).where(ResourceLock.run_id == run.id))


async def reclaim_expired_leases(session: AsyncSession, host: TestHost) -> int:
    """Requeue safe work and freeze ambiguous non-idempotent work."""

    now = utcnow()
    expired = (
        await session.scalars(
            select(RunStep)
            .join(Run, Run.id == RunStep.run_id)
            .where(
                Run.test_host_id == host.id,
                Run.state.in_([RunState.RUNNING.value, RunState.RECOVERING.value]),
                RunStep.state.in_([StepState.DISPATCHED.value, StepState.RUNNING.value]),
                RunStep.lease_expires_at.is_not(None),
                RunStep.lease_expires_at <= now,
            )
            .order_by(RunStep.run_id, RunStep.seq)
            .with_for_update()
        )
    ).all()

    for step in expired:
        run = await _locked_run(session, step.run_id)
        previous = step.state
        step.leased_by = None
        step.lease_expires_at = None

        if previous == StepState.DISPATCHED.value:
            step.state = StepState.PENDING.value
            await append_event(
                session,
                run,
                EventKind.STEP_LEASE_EXPIRED.value,
                "Agent 未 ACK，步骤租约已回收",
                severity=Severity.WARNING.value,
                payload={"attempt": step.attempt, "previous_state": previous},
                step_id=step.id,
                loop_index=step.loop_index,
            )
            continue

        can_retry = step.idempotent and step.attempt <= step.max_retry
        if can_retry:
            step.state = StepState.PENDING.value
            await append_event(
                session,
                run,
                EventKind.STEP_RETRYING.value,
                "Agent 执行中失联，幂等步骤已安全回收等待重试",
                severity=Severity.WARNING.value,
                payload={"attempt": step.attempt, "max_retry": step.max_retry},
                step_id=step.id,
                loop_index=step.loop_index,
            )
            continue

        step.state = StepState.TIMED_OUT.value
        step.ended_at = now
        await append_event(
            session,
            run,
            EventKind.STEP_TIMED_OUT.value,
            "执行租约到期，步骤结果不确定",
            severity=Severity.ERROR.value,
            payload={
                "attempt": step.attempt,
                "idempotent": step.idempotent,
                "retry_exhausted": step.idempotent,
            },
            step_id=step.id,
            loop_index=step.loop_index,
        )
        await _freeze_run(
            session,
            run,
            step,
            "步骤 {0} 执行结果不确定，已冻结现场禁止盲目重跑".format(step.step_key),
        )

    return len(expired)


async def try_lease_next_step(
    session: AsyncSession, host: TestHost
) -> Optional[Dict[str, Any]]:
    await reclaim_expired_leases(session, host)
    runs = (
        await session.scalars(
            select(Run)
            .where(
                Run.test_host_id == host.id,
                Run.state == RunState.RUNNING.value,
            )
            .order_by(Run.created_at, Run.id)
            .with_for_update(skip_locked=True)
        )
    ).all()

    for run in runs:
        steps = (
            await session.scalars(
                select(RunStep)
                .where(RunStep.run_id == run.id, RunStep.state.in_(_ACTIVE_STEP_STATES))
                .order_by(RunStep.seq)
                .with_for_update(skip_locked=True)
            )
        ).all()
        if not steps:
            continue
        step = steps[0]
        if step.state != StepState.PENDING.value:
            continue

        step.state = StepState.DISPATCHED.value
        step.attempt += 1
        step.leased_by = host.agent_id
        step.lease_expires_at = utcnow() + _dt.timedelta(
            seconds=get_settings().timing.lease_ttl_s
        )
        await append_event(
            session,
            run,
            EventKind.STEP_DISPATCHED.value,
            "步骤已租给 Host Agent，等待 ACK",
            payload={"agent_id": host.agent_id, "attempt": step.attempt},
            step_id=step.id,
            loop_index=step.loop_index,
        )
        return step_lease(step, run)
    return None


async def acknowledge_step(
    session: AsyncSession,
    host: TestHost,
    step_id: str,
    attempt: int,
    idempotency_key: str,
) -> Dict[str, Any]:
    step, run = await _owned_step(session, host, step_id)
    if await _replayed(
        session,
        run_id=run.id,
        step_id=step.id,
        idempotency_key=idempotency_key,
        operation="ack",
    ):
        return step_status(step, replayed=True)
    if step.state != StepState.DISPATCHED.value:
        raise StepLeaseConflict("只有 DISPATCHED 步骤可以 ACK")
    _assert_current(step, host.agent_id or "", attempt)

    now = utcnow()
    step.state = StepState.RUNNING.value
    step.started_at = now
    step.lease_expires_at = min(
        now + _dt.timedelta(seconds=get_settings().timing.lease_ttl_s),
        now + _dt.timedelta(seconds=step.timeout_s),
    )
    await append_event(
        session,
        run,
        EventKind.STEP_STARTED.value,
        "Host Agent 已确认接手步骤",
        source=EventSource.AGENT.value,
        payload={"_operation": "ack", "attempt": attempt},
        idempotency_key=idempotency_key,
        step_id=step.id,
        loop_index=step.loop_index,
    )
    return step_status(step)

async def renew_step(
    session: AsyncSession,
    host: TestHost,
    step_id: str,
    attempt: int,
    idempotency_key: str,
) -> Dict[str, Any]:
    step, run = await _owned_step(session, host, step_id)
    if await _replayed(
        session,
        run_id=run.id,
        step_id=step.id,
        idempotency_key=idempotency_key,
        operation="renew",
    ):
        return step_status(step, replayed=True)
    if step.state != StepState.RUNNING.value:
        raise StepLeaseConflict("只有 RUNNING 步骤可以续租")
    _assert_current(step, host.agent_id or "", attempt)
    assert step.started_at is not None
    now = utcnow()
    step.lease_expires_at = min(
        now + _dt.timedelta(seconds=get_settings().timing.lease_ttl_s),
        step.started_at + _dt.timedelta(seconds=step.timeout_s),
    )
    await append_event(
        session,
        run,
        EventKind.STEP_LEASE_RENEWED.value,
        "Host Agent 已续租执行步骤",
        source=EventSource.AGENT.value,
        severity=Severity.DEBUG.value,
        payload={"_operation": "renew", "attempt": attempt},
        idempotency_key=idempotency_key,
        step_id=step.id,
        loop_index=step.loop_index,
    )
    return step_status(step)


async def record_step_events(
    session: AsyncSession,
    host: TestHost,
    step_id: str,
    attempt: int,
    events: List[AgentStepEvent],
    idempotency_key: str,
) -> Dict[str, Any]:
    step, run = await _owned_step(session, host, step_id)
    first_key = "{0}:0".format(idempotency_key)
    if await _replayed(
        session,
        run_id=run.id,
        step_id=step.id,
        idempotency_key=first_key,
        operation="events",
    ):
        return {**step_status(step, replayed=True), "accepted": len(events)}
    if step.state != StepState.RUNNING.value:
        raise StepLeaseConflict("只有 RUNNING 步骤可以上报事件")
    _assert_current(step, host.agent_id or "", attempt)

    for index, item in enumerate(events):
        await append_event(
            session,
            run,
            EventKind.STEP_OUTPUT.value,
            item.message,
            source=EventSource.AGENT.value,
            severity=item.severity,
            payload={"_operation": "events", "attempt": attempt, **item.payload},
            idempotency_key="{0}:{1}".format(idempotency_key, index),
            step_id=step.id,
            loop_index=step.loop_index,
        )
    return {**step_status(step), "accepted": len(events)}


async def _finish_run(session: AsyncSession, run: Run, verdict: str) -> None:
    await append_event(
        session,
        run,
        EventKind.LOCK_RELEASED.value,
        "DUT 与测试主机独占锁已释放",
    )
    await session.execute(delete(ResourceLock).where(ResourceLock.run_id == run.id))
    await append_event(
        session,
        run,
        EventKind.RUN_COMPLETED.value,
        "独立 Host Agent 已完成全部步骤，结论 {0}".format(verdict),
        target_state=RunState.COMPLETED.value,
        projection={"verdict": verdict, "ended_at": utcnow()},
    )
    bundle = await build_evidence_bundle(session, run)
    await append_event(
        session,
        run,
        EventKind.EVIDENCE_BUNDLED.value,
        "Evidence Bundle 已生成，完整率 100%",
        payload={"bundle_id": bundle.id, "uri": bundle.uri, "completeness": 1.0},
    )


async def complete_step(
    session: AsyncSession,
    host: TestHost,
    step_id: str,
    body: AgentStepCompleteBody,
    idempotency_key: str,
) -> Dict[str, Any]:
    step, run = await _owned_step(session, host, step_id)
    if await _replayed(
        session,
        run_id=run.id,
        step_id=step.id,
        idempotency_key=idempotency_key,
        operation="complete",
    ):
        return step_status(step, replayed=True)
    if step.state != StepState.RUNNING.value:
        raise StepLeaseConflict("只有 RUNNING 步骤可以完成")
    _assert_current(step, host.agent_id or "", body.attempt)

    step.exit_code = body.exit_code
    step.error_class = body.error_class
    step.result = body.result
    step.log_uri = body.log_uri
    step.ended_at = utcnow()
    step.leased_by = None
    step.lease_expires_at = None

    succeeded = body.status == StepState.SUCCEEDED.value
    step.state = StepState.SUCCEEDED.value if succeeded else StepState.FAILED.value
    await append_event(
        session,
        run,
        EventKind.STEP_SUCCEEDED.value if succeeded else EventKind.STEP_FAILED.value,
        "步骤执行成功" if succeeded else "步骤执行失败",
        source=EventSource.AGENT.value,
        severity=Severity.INFO.value if succeeded else Severity.ERROR.value,
        payload={
            "_operation": "complete",
            "attempt": body.attempt,
            "exit_code": body.exit_code,
            "error_class": body.error_class,
            "result": body.result,
            "log_uri": body.log_uri,
        },
        idempotency_key=idempotency_key,
        step_id=step.id,
        loop_index=step.loop_index,
    )

    if not succeeded:
        if (
            step.on_fail == OnFail.RETRY.value
            and step.idempotent
            and step.attempt <= step.max_retry
        ):
            step.state = StepState.PENDING.value
            await append_event(
                session,
                run,
                EventKind.STEP_RETRYING.value,
                "幂等步骤失败，已进入受限重试队列",
                severity=Severity.WARNING.value,
                payload={"attempt": step.attempt, "max_retry": step.max_retry},
                step_id=step.id,
                loop_index=step.loop_index,
            )
            return step_status(step)
        if step.on_fail in {OnFail.FREEZE.value, OnFail.RECOVER.value} or not step.idempotent:
            await _freeze_run(
                session,
                run,
                step,
                "步骤 {0} 失败且不可安全自动重跑，已冻结现场".format(step.step_key),
            )
            return step_status(step)
        if step.on_fail == OnFail.FAIL_RUN.value:
            await _finish_run(session, run, Verdict.FAIL.value)
            return step_status(step)

    remaining = await session.scalar(
        select(func.count())
        .select_from(RunStep)
        .where(RunStep.run_id == run.id, RunStep.state.in_(_ACTIVE_STEP_STATES))
    )

    if succeeded and step.checkpoint_after:
        next_step = await session.scalar(
            select(RunStep)
            .where(RunStep.run_id == run.id, RunStep.seq > step.seq)
            .order_by(RunStep.seq)
            .limit(1)
        )
        checkpoint = {
            "loop_index": step.loop_index,
            "completed_step_key": step.step_key,
            "next_step_key": next_step.step_key if next_step else None,
        }
        loop_done = max(run.loop_done, step.loop_index or run.loop_done)
        await append_event(
            session,
            run,
            EventKind.CHECKPOINT_SAVED.value,
            "安全检查点已保存",
            payload=checkpoint,
            projection={"checkpoint": checkpoint, "loop_done": loop_done},
            step_id=step.id,
            loop_index=step.loop_index,
        )

    if not remaining:
        failed = await session.scalar(
            select(func.count())
            .select_from(RunStep)
            .where(RunStep.run_id == run.id, RunStep.state == StepState.FAILED.value)
        )
        await _finish_run(
            session,
            run,
            Verdict.FAIL.value if failed else Verdict.PASS.value,
        )
    return step_status(step)
