from __future__ import annotations

from collections.abc import Mapping
from datetime import UTC, datetime, timedelta

from sqlalchemy import select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from backup_tool.db.models import Execution, ExecutionEvent, Job
from backup_tool.security.redaction import redact

ACTIVE_STATES = frozenset({"queued", "preparing", "running", "verifying", "cancelling"})
TERMINAL_STATES = frozenset({"committed", "cancelled", "failed"})
TRANSIENT_REASONS = frozenset({"worker_lost", "timeout", "unavailable", "transient_io"})

_ALLOWED: Mapping[str, frozenset[str]] = {
    "queued": frozenset({"preparing", "cancelled", "failed"}),
    "preparing": frozenset({"running", "cancelling", "failed"}),
    "running": frozenset({"verifying", "cancelling", "failed"}),
    "verifying": frozenset({"committed", "failed"}),
    "cancelling": frozenset({"cancelled", "failed"}),
    "committed": frozenset(),
    "cancelled": frozenset(),
    "failed": frozenset(),
}


class TransitionError(ValueError):
    code = "invalid_transition"


class EnqueueError(ValueError):
    def __init__(self, code: str, detail: str, active_execution_id: str | None = None):
        super().__init__(detail)
        self.code = code
        self.active_execution_id = active_execution_id


def transition(current: str, target: str) -> str:
    """Validate and return one durable execution-state transition."""
    if target not in _ALLOWED.get(current, frozenset()):
        raise TransitionError(f"Cannot transition execution from {current!r} to {target!r}.")
    return target


async def enqueue(db: AsyncSession, job_id: str, trigger: str = "manual") -> Execution:
    """Create exactly one active execution for an enabled active job."""
    job = await db.get(Job, job_id)
    if job is None:
        raise EnqueueError("resource_not_found", "Job was not found.")
    if job.state != "active" or not job.enabled:
        raise EnqueueError("job_disabled", "Job is disabled or archived.")
    job_identifier = job.id
    execution = Execution(job_id=job_identifier, trigger=trigger, progress={})
    db.add(execution)
    await db.flush()
    await record_event(db, execution)
    try:
        await db.commit()
    except IntegrityError as error:
        await db.rollback()
        existing = await db.scalar(
            select(Execution.id).where(
                Execution.job_id == job_identifier, Execution.state.in_(ACTIVE_STATES)
            )
        )
        raise EnqueueError(
            "execution_active", "Job already has an active execution.", existing
        ) from error
    await db.refresh(execution)
    await record_event(db, execution)
    await db.commit()
    return execution


async def claim(
    db: AsyncSession, execution_id: str, owner: str, lease_seconds: int = 60
) -> Execution | None:
    """Atomically lease a queued execution; a stale lease may be recovered."""
    now = datetime.now(UTC)
    expires = now + timedelta(seconds=lease_seconds)
    result = await db.execute(
        update(Execution)
        .where(
            Execution.id == execution_id,
            Execution.state.in_({"queued", "preparing"}),
            (Execution.lease_expires_at.is_(None)) | (Execution.lease_expires_at < now),
        )
        .values(state="preparing", lease_owner=owner, lease_expires_at=expires, heartbeat_at=now)
    )
    if getattr(result, "rowcount", 0) != 1:
        await db.rollback()
        return None
    await db.commit()
    execution = await db.get(Execution, execution_id)
    if execution is not None:
        await record_event(db, execution)
        await db.commit()
    return execution


async def heartbeat(
    db: AsyncSession, execution_id: str, owner: str, lease_seconds: int = 60
) -> bool:
    now = datetime.now(UTC)
    result = await db.execute(
        update(Execution)
        .where(
            Execution.id == execution_id,
            Execution.lease_owner == owner,
            Execution.state.in_(ACTIVE_STATES),
            Execution.lease_expires_at >= now,
        )
        .values(heartbeat_at=now, lease_expires_at=now + timedelta(seconds=lease_seconds))
    )
    await db.commit()
    return getattr(result, "rowcount", 0) == 1


async def request_cancellation(db: AsyncSession, execution_id: str) -> Execution | None:
    # Queued work has no worker and may terminate immediately; leased work asks its owner.
    now = datetime.now(UTC)
    result = await db.execute(
        update(Execution)
        .where(Execution.id == execution_id, Execution.state == "queued")
        .values(state="cancelled", reason_code="cancellation_requested", completed_at=now)
    )
    if getattr(result, "rowcount", 0) != 1:
        result = await db.execute(
            update(Execution)
            .where(Execution.id == execution_id, Execution.state.in_({"preparing", "running"}))
            .values(state="cancelling", reason_code="cancellation_requested")
        )
    if getattr(result, "rowcount", 0) != 1:
        await db.rollback()
        return None
    await db.commit()
    execution = await db.get(Execution, execution_id)
    if execution is not None:
        await record_event(db, execution)
        await db.commit()
    return execution


async def complete_cancellation(db: AsyncSession, execution_id: str, owner: str) -> bool:
    result = await db.execute(
        update(Execution)
        .where(
            Execution.id == execution_id,
            Execution.lease_owner == owner,
            Execution.lease_expires_at >= datetime.now(UTC),
            Execution.state == "cancelling",
        )
        .values(
            state="cancelled",
            completed_at=datetime.now(UTC),
            lease_owner=None,
            lease_expires_at=None,
        )
    )
    await db.commit()
    return getattr(result, "rowcount", 0) == 1


async def retry(db: AsyncSession, execution_id: str) -> Execution | None:
    execution = await db.get(Execution, execution_id)
    if execution is None:
        return None
    if execution.state != "failed" or execution.reason_code not in TRANSIENT_REASONS:
        raise EnqueueError("retry_not_allowed", "Execution failure is not retryable.")
    execution.state = "queued"
    execution.attempt += 1
    execution.lease_owner = None
    execution.lease_expires_at = None
    execution.heartbeat_at = None
    execution.reason_code = None
    execution.operator_message = None
    try:
        await db.commit()
    except IntegrityError as error:
        await db.rollback()
        active_id = await db.scalar(
            select(Execution.id).where(
                Execution.job_id == execution.job_id, Execution.state.in_(ACTIVE_STATES)
            )
        )
        raise EnqueueError(
            "execution_active", "Job already has an active execution.", active_id
        ) from error
    await db.refresh(execution)
    await record_event(db, execution)
    await db.commit()
    return execution


def _redact_progress(value: object) -> object:
    """Apply common redaction plus execution-specific path redaction recursively."""
    if isinstance(value, Mapping):
        safe: dict[str, object] = {}
        for key, item in value.items():
            normalized = str(key).lower()
            if any(marker in normalized for marker in ("path", "secret", "token", "password")):
                safe[str(key)] = "[REDACTED]"
            else:
                safe[str(key)] = _redact_progress(item)
        return safe
    if isinstance(value, list):
        return [_redact_progress(item) for item in value]
    return redact(value)


async def record_event(db: AsyncSession, execution: Execution) -> ExecutionEvent:
    """Append the public representation as the next durable event revision."""
    execution.revision += 1
    event = ExecutionEvent(
        execution_id=execution.id,
        revision=execution.revision,
        payload=public_event(execution),
    )
    db.add(event)
    await db.flush()
    return event


def public_event(execution: Execution) -> dict[str, object]:
    """Return redacted progress suitable for polling or SSE."""
    progress = _redact_progress(execution.progress)
    return {
        "id": execution.id,
        "state": execution.state,
        "attempt": execution.attempt,
        "revision": execution.revision,
        "reason_code": execution.reason_code,
        "progress": progress,
    }


async def recover_stale(db: AsyncSession) -> int:
    now = datetime.now(UTC)
    cancelled = await db.execute(
        update(Execution)
        .where(Execution.state == "cancelling", Execution.lease_expires_at < now)
        .values(
            state="cancelled",
            completed_at=now,
            lease_owner=None,
            lease_expires_at=None,
            heartbeat_at=None,
            reason_code="cancellation_requested",
        )
    )
    recovered = await db.execute(
        update(Execution)
        .where(
            Execution.state.in_({"preparing", "running", "verifying"}),
            Execution.lease_expires_at < now,
        )
        .values(
            state="queued",
            lease_owner=None,
            lease_expires_at=None,
            heartbeat_at=None,
            reason_code="worker_lost",
        )
    )
    await db.commit()
    return (getattr(cancelled, "rowcount", 0) or 0) + (getattr(recovered, "rowcount", 0) or 0)