from __future__ import annotations

from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
import logging
import socket

from sqlalchemy import func, or_, select
from sqlalchemy.orm import Session

from govoplan_core.auth import ApiPrincipal
from govoplan_core.core.automation import (
    AutomationPrincipalRequest,
    automation_principal_provider,
)
from govoplan_core.core.notifications import (
    NotificationDispatchRequest,
    notification_dispatch_provider,
)
from govoplan_core.db.base import utcnow
from govoplan_dataflow.backend.db.models import (
    DataflowPipeline,
    DataflowPipelineRevision,
    DataflowRun,
    DataflowTrigger,
)
from govoplan_dataflow.backend.governance import require_definition_action
from govoplan_dataflow.backend.recovery import (
    DataflowRecoveryError,
    DataflowRunRecovery,
    begin_dataflow_run_recovery,
    claim_stale_dataflow_recovery,
    dataflow_run_recovery_state,
)
from govoplan_dataflow.backend.service import (
    _execute_pipeline_run,
    _require_deployed_revision,
    pipeline_run_request,
)


MAX_ACTIVE_RUNS_PER_TENANT = 4
DEFAULT_LEASE_SECONDS = 120
logger = logging.getLogger(__name__)


class DataflowWorkerError(RuntimeError):
    pass


def dispatch_pending_runs(
    session: Session,
    *,
    registry: object | None,
    now: datetime | None = None,
    limit: int = 10,
    worker_id: str | None = None,
) -> dict[str, object]:
    current = _as_utc(now or utcnow())
    resolved_worker_id = (
        str(worker_id).strip()
        if worker_id and str(worker_id).strip()
        else socket.gethostname()
    )
    recovered, recovered_outcome_unknown = _recover_expired_leases(
        session,
        now=current,
    )
    session.commit()
    summary: dict[str, object] = {
        "claimed": 0,
        "succeeded": 0,
        "retrying": 0,
        "failed": 0,
        "cancelled": 0,
        "outcome_unknown": recovered_outcome_unknown,
        "recovered": recovered,
        "runs": [],
    }
    for _ in range(max(1, min(int(limit), 100))):
        run_id = _claim_next_run(
            session,
            now=current,
            worker_id=resolved_worker_id,
        )
        if run_id is None:
            session.rollback()
            break
        session.commit()
        summary["claimed"] = int(summary["claimed"]) + 1
        outcome = _execute_claimed_run(
            session,
            run_id=run_id,
            registry=registry,
            now=current,
        )
        session.commit()
        summary[outcome] = int(summary[outcome]) + 1
        runs = list(summary["runs"])
        runs.append({"ref": f"dataflow-run:{run_id}", "status": outcome})
        summary["runs"] = runs
    return summary


def purge_expired_runs(
    session: Session,
    *,
    now: datetime | None = None,
    limit: int = 500,
) -> dict[str, object]:
    current = _as_utc(now or utcnow())
    runs = list(
        session.scalars(
            select(DataflowRun)
            .where(
                DataflowRun.status.in_(
                    ("succeeded", "failed", "cancelled")
                ),
                DataflowRun.retention_until.is_not(None),
                DataflowRun.retention_until <= current,
                DataflowRun.purged_at.is_(None),
            )
            .order_by(DataflowRun.retention_until, DataflowRun.id)
            .limit(max(1, min(int(limit), 5_000)))
            .with_for_update(skip_locked=True)
        )
    )
    for run in runs:
        run.request_ = {}
        run.source_fingerprints = []
        run.diagnostics = []
        authorization = dict(run.authorization_)
        authorization.pop("submitted_principal", None)
        authorization["personal_data_purged"] = True
        run.authorization_ = authorization
        run.purged_at = current
    session.flush()
    return {
        "purged": len(runs),
        "run_refs": [f"dataflow-run:{run.id}" for run in runs],
    }


def run_metrics(
    session: Session,
    *,
    tenant_id: str,
    now: datetime | None = None,
) -> dict[str, object]:
    current = _as_utc(now or utcnow())
    rows = session.execute(
        select(DataflowRun.status, func.count())
        .where(DataflowRun.tenant_id == tenant_id)
        .group_by(DataflowRun.status)
    ).all()
    statuses = {str(status): int(count) for status, count in rows}
    recent = list(
        session.execute(
            select(
                DataflowRun.status,
                DataflowRun.started_at,
                DataflowRun.finished_at,
            ).where(
                DataflowRun.tenant_id == tenant_id,
                DataflowRun.finished_at.is_not(None),
                DataflowRun.finished_at >= current - timedelta(hours=24),
            )
        )
    )
    durations = [
        (
            _as_utc(finished_at) - _as_utc(started_at)
        ).total_seconds()
        for _status, started_at, finished_at in recent
        if finished_at is not None and started_at is not None
    ]
    oldest_queued_at = session.scalar(
        select(func.min(DataflowRun.queued_at)).where(
            DataflowRun.tenant_id == tenant_id,
            DataflowRun.status.in_(("queued", "retrying")),
        )
    )
    return {
        "statuses": statuses,
        "active": sum(
            statuses.get(status, 0)
            for status in ("queued", "retrying", "running")
        ),
        "completed_last_24_hours": len(recent),
        "failed_last_24_hours": sum(
            1 for status, _started_at, _finished_at in recent
            if status == "failed"
        ),
        "average_duration_seconds": (
            round(sum(durations) / len(durations), 3)
            if durations
            else None
        ),
        "oldest_queued_at": oldest_queued_at,
    }


def _recover_expired_leases(
    session: Session,
    *,
    now: datetime,
) -> tuple[int, int]:
    runs = list(
        session.scalars(
            select(DataflowRun)
            .where(
                DataflowRun.status == "running",
                DataflowRun.lease_expires_at.is_not(None),
                DataflowRun.lease_expires_at < now,
            )
            .with_for_update(skip_locked=True)
        )
    )
    recovered = 0
    outcome_unknown = 0
    for run in runs:
        recovery = dataflow_run_recovery_state(session, run_id=run.id)
        recovery_status: str | None = None
        recovery_metadata = dict(run.authorization_).get("recovery")
        effect_started = bool(
            isinstance(recovery_metadata, Mapping)
            and recovery_metadata.get("boundary") == "output-publication"
        )
        if recovery is not None and recovery.get("operation_id"):
            try:
                recovery_status = claim_stale_dataflow_recovery(
                    session,
                    operation_id=str(recovery["operation_id"]),
                    lease_ttl_seconds=DEFAULT_LEASE_SECONDS,
                    effect_started=effect_started,
                )
            except DataflowRecoveryError:
                logger.warning(
                    "Could not claim stale Dataflow recovery operation %s",
                    recovery["operation_id"],
                    exc_info=True,
                )
                continue
        run.worker_id = None
        run.claimed_at = None
        run.lease_expires_at = None
        run.heartbeat_at = None
        if run.cancellation_requested_at is not None:
            _cancel_run(run, now=now)
        elif recovery_status == "outcome_unknown" or effect_started:
            run.status = "outcome_unknown"
            run.finished_at = now
            run.progress_phase = "outcome_unknown"
            run.error = (
                "The worker lease expired after output publication may have "
                "started. Verify the sink before retrying."
            )
            outcome_unknown += 1
        elif run.attempts < run.max_attempts:
            run.status = "retrying"
            run.available_at = now
            run.progress_phase = "retrying_after_worker_loss"
            run.error = "The previous worker lease expired."
        else:
            _fail_run(
                run,
                now=now,
                message="The worker lease expired after the final attempt.",
            )
        recovered += 1
    return recovered, outcome_unknown


def _claim_next_run(
    session: Session,
    *,
    now: datetime,
    worker_id: str,
) -> str | None:
    candidates = list(
        session.scalars(
            select(DataflowRun)
            .where(
                DataflowRun.status.in_(("queued", "retrying")),
                or_(
                    DataflowRun.available_at.is_(None),
                    DataflowRun.available_at <= now,
                ),
            )
            .order_by(DataflowRun.available_at, DataflowRun.created_at)
            .limit(40)
            .with_for_update(skip_locked=True)
        )
    )
    for run in candidates:
        active = int(
            session.scalar(
                select(func.count())
                .select_from(DataflowRun)
                .where(
                    DataflowRun.tenant_id == run.tenant_id,
                    DataflowRun.status == "running",
                )
            )
            or 0
        )
        if active >= MAX_ACTIVE_RUNS_PER_TENANT:
            continue
        run.status = "running"
        run.attempts += 1
        run.worker_id = worker_id[:255]
        run.claimed_at = now
        run.heartbeat_at = now
        run.lease_expires_at = now + timedelta(
            seconds=_lease_seconds(run)
        )
        run.started_at = run.started_at or now
        run.finished_at = None
        run.progress_percent = max(run.progress_percent, 5)
        run.progress_phase = "authorizing"
        session.flush()
        return run.id
    return None


def _execute_claimed_run(
    session: Session,
    *,
    run_id: str,
    registry: object | None,
    now: datetime,
) -> str:
    run = session.get(DataflowRun, run_id)
    if run is None:
        return "failed"
    if run.cancellation_requested_at is not None:
        _cancel_run(run, now=now)
        return "cancelled"
    recovery: DataflowRunRecovery | None = None
    try:
        principal, provenance = _resolve_principal(
            session,
            run=run,
            registry=registry,
        )
        run.authorization_ = {
            **dict(run.authorization_),
            "last_resolution": provenance,
            "resolved_at": now.isoformat(),
        }
        pipeline = session.get(DataflowPipeline, run.pipeline_id)
        revision = session.get(
            DataflowPipelineRevision,
            run.pipeline_revision_id,
        )
        if pipeline is None or revision is None:
            raise DataflowWorkerError(
                "The pipeline or pinned revision no longer exists."
            )
        action = (
            "run"
            if run.invocation_kind in {"manual", "api", "backfill"}
            else "automate"
        )
        require_definition_action(
            pipeline,
            principal=principal,
            registry=registry,
            action=action,
        )
        _require_deployed_revision(
            session,
            tenant_id=run.tenant_id,
            pipeline=pipeline,
            revision=revision,
            environment=run.environment,
        )
        session.commit()
        recovery = begin_dataflow_run_recovery(
            session,
            run=run,
            lease_ttl_seconds=_lease_seconds(run),
        )
        if recovery.replayed:
            session.rollback()
            run = session.get(DataflowRun, run_id)
            return run.status if run is not None else "failed"
        _notify_run(
            session,
            registry=registry,
            run=run,
            event_kind="dataflow.run.started",
            subject=f"Dataflow run started: {pipeline.name}",
        )
        retryable = _execute_pipeline_run(
            session,
            run=run,
            pipeline=pipeline,
            revision=revision,
            request=pipeline_run_request(run),
            principal=principal,
            registry=registry,
            recovery=recovery,
        )
        if (
            run.status == "failed"
            and retryable
            and run.attempts < run.max_attempts
        ):
            _retry_run(run, now=now)
        _finish_claim(run)
        _update_trigger_status(session, run)
        _notify_terminal_run(
            session,
            registry=registry,
            run=run,
            pipeline_name=pipeline.name,
        )
        recovery.finish(session, run=run)
        return run.status
    except DataflowRecoveryError:
        logger.exception(
            "Dataflow recovery finalization failed for run %s",
            run_id,
        )
        session.rollback()
        run = session.get(DataflowRun, run_id)
        if run is not None:
            if recovery is not None and recovery.publication_started:
                run.status = "outcome_unknown"
                run.finished_at = now
                run.progress_phase = "outcome_unknown"
                run.error = (
                    "Output publication completed without verifiable terminal "
                    "recovery evidence; reconcile the sink."
                )
                _finish_claim(run)
                return "outcome_unknown"
            if recovery is None:
                _fail_run(
                    run,
                    now=now,
                    message=(
                        "The recovery ledger was unavailable; no Dataflow "
                        "execution started."
                    ),
                )
                _finish_claim(run)
            else:
                _fail_run(
                    run,
                    now=now,
                    message=(
                        "The database-only run could not commit verifiable "
                        "recovery evidence."
                    ),
                )
                _finish_claim(run)
        return "failed"
    except (DataflowWorkerError, PermissionError, ValueError) as exc:
        _fail_run(run, now=now, message=str(exc))
        _finish_claim(run)
        _update_trigger_status(session, run)
        _notify_terminal_run(
            session,
            registry=registry,
            run=run,
            pipeline_name="Dataflow",
        )
        if recovery is not None:
            recovery.finish(session, run=run)
        return "failed"
    except Exception as exc:
        logger.exception(
            "Unexpected Dataflow worker failure for run %s",
            run_id,
        )
        session.rollback()
        run = session.get(DataflowRun, run_id)
        if run is None:
            return "failed"
        if recovery is not None and recovery.publication_started:
            run.status = "outcome_unknown"
            run.finished_at = now
            run.error = (
                "The output provider failed after dispatch began; verify the "
                "sink before retrying."
            )
            run.progress_phase = "outcome_unknown"
            _finish_claim(run)
            try:
                recovery.finish(session, run=run)
            except DataflowRecoveryError:
                logger.exception(
                    "Could not persist the uncertain output state for run %s",
                    run_id,
                )
                session.rollback()
            return "outcome_unknown"
        _fail_run(
            run,
            now=now,
            message=(
                "Unexpected worker failure "
                f"({type(exc).__name__})."
            ),
        )
        _finish_claim(run)
        if recovery is not None:
            try:
                recovery.finish(session, run=run)
            except DataflowRecoveryError:
                logger.exception(
                    "Could not finalize failed Dataflow recovery for run %s",
                    run_id,
                )
                session.rollback()
        return "failed"


def _resolve_principal(
    session: Session,
    *,
    run: DataflowRun,
    registry: object | None,
) -> tuple[ApiPrincipal, Mapping[str, object]]:
    provider = automation_principal_provider(registry)
    if provider is None:
        raise DataflowWorkerError(
            "Queued Dataflow execution requires the automation-principal "
            "capability."
        )
    authorization = dict(run.authorization_)
    subject_kind = str(
        authorization.get("subject_kind") or "delegated_user"
    )
    common = {
        "tenant_id": run.tenant_id,
        "authorization_ref": str(
            authorization.get("authorization_ref")
            or f"dataflow-run:{run.id}"
        ),
        "grant_scopes": tuple(
            str(scope)
            for scope in authorization.get("grant_scopes") or ()
        ),
        "context": {
            "run_ref": f"dataflow-run:{run.id}",
            "pipeline_ref": f"pipeline:{run.pipeline_id}",
            "environment": run.environment,
            "attempt": run.attempts,
        },
    }
    if subject_kind == "service_account":
        service_account_id = str(
            authorization.get("service_account_id") or ""
        ).strip()
        request = AutomationPrincipalRequest.service_account(
            service_account_id=service_account_id,
            **common,
        )
    else:
        request = AutomationPrincipalRequest.delegated_user(
            account_id=str(authorization.get("account_id") or ""),
            membership_id=str(
                authorization.get("membership_id") or ""
            ),
            **common,
        )
    resolution = provider.resolve_automation_principal(
        session,
        request=request,
    )
    if not resolution.allowed or not isinstance(
        resolution.principal,
        ApiPrincipal,
    ):
        raise DataflowWorkerError(
            resolution.reason or "Dataflow run authorization was denied."
        )
    return resolution.principal, resolution.provenance


def _retry_run(run: DataflowRun, *, now: datetime) -> None:
    delay_seconds = min(15 * (2 ** max(0, run.attempts - 1)), 900)
    run.status = "retrying"
    run.available_at = now + timedelta(seconds=delay_seconds)
    run.finished_at = None
    run.progress_percent = 0
    run.progress_phase = "retrying"


def _finish_claim(run: DataflowRun) -> None:
    run.worker_id = None
    run.claimed_at = None
    run.lease_expires_at = None
    run.heartbeat_at = None


def _cancel_run(run: DataflowRun, *, now: datetime) -> None:
    run.status = "cancelled"
    run.finished_at = now
    run.error = "Cancelled by request."
    run.progress_phase = "cancelled"
    _finish_claim(run)


def _fail_run(
    run: DataflowRun,
    *,
    now: datetime,
    message: str,
) -> None:
    run.status = "failed"
    run.finished_at = now
    run.error = message
    run.progress_phase = "failed"


def _update_trigger_status(session: Session, run: DataflowRun) -> None:
    if not run.trigger_id:
        return
    trigger = session.get(DataflowTrigger, run.trigger_id)
    if trigger is None:
        return
    trigger.last_status = run.status
    trigger.last_error = run.error


def _notify_terminal_run(
    session: Session,
    *,
    registry: object | None,
    run: DataflowRun,
    pipeline_name: str,
) -> None:
    if run.status == "retrying":
        event_kind = "dataflow.run.retrying"
        subject = f"Dataflow run will retry: {pipeline_name}"
    elif run.status == "succeeded":
        event_kind = "dataflow.run.completed"
        subject = f"Dataflow run completed: {pipeline_name}"
    elif run.status == "cancelled":
        event_kind = "dataflow.run.cancelled"
        subject = f"Dataflow run cancelled: {pipeline_name}"
    else:
        event_kind = "dataflow.run.failed"
        subject = f"Dataflow run failed: {pipeline_name}"
    _notify_run(
        session,
        registry=registry,
        run=run,
        event_kind=event_kind,
        subject=subject,
    )


def _notify_run(
    session: Session,
    *,
    registry: object | None,
    run: DataflowRun,
    event_kind: str,
    subject: str,
) -> None:
    provider = notification_dispatch_provider(registry)
    authorization = dict(run.authorization_)
    recipient_id = str(authorization.get("account_id") or "").strip()
    if provider is None or not recipient_id:
        return
    try:
        provider.enqueue_notification(
            session,
            NotificationDispatchRequest(
                tenant_id=run.tenant_id,
                source_module="dataflow",
                source_resource_type="dataflow_run",
                source_resource_id=run.id,
                event_kind=event_kind,
                recipient_type="account",
                recipient_id=recipient_id,
                subject=subject,
                body_text=run.error,
                action_url="/dataflow",
                payload={
                    "run_ref": f"dataflow-run:{run.id}",
                    "pipeline_ref": f"pipeline:{run.pipeline_id}",
                    "status": run.status,
                    "environment": run.environment,
                    "attempt": run.attempts,
                },
            ),
        )
    except Exception:
        # Notification delivery is optional and cannot change run outcome.
        logger.warning(
            "Could not enqueue Dataflow notification for run %s",
            run.id,
            exc_info=True,
        )
        return


def _lease_seconds(run: DataflowRun) -> int:
    budget = dict(run.resource_budget)
    wall_seconds = float(budget.get("max_wall_seconds") or 30.0)
    return max(DEFAULT_LEASE_SECONDS, int(wall_seconds) + 60)


def _as_utc(value: datetime | None) -> datetime:
    if value is None:
        return datetime.now(timezone.utc)
    if value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


class SqlDataflowRunWorker:
    def __init__(self, *, registry: object | None = None) -> None:
        self._registry = registry

    def dispatch_pending(
        self,
        session: object,
        *,
        now: datetime | None = None,
        limit: int = 10,
        worker_id: str | None = None,
    ) -> Mapping[str, object]:
        if not isinstance(session, Session):
            raise TypeError("Dataflow run dispatch requires a Session.")
        return dispatch_pending_runs(
            session,
            registry=self._registry,
            now=now,
            limit=limit,
            worker_id=worker_id,
        )

    def purge_expired(
        self,
        session: object,
        *,
        now: datetime | None = None,
        limit: int = 500,
    ) -> Mapping[str, object]:
        if not isinstance(session, Session):
            raise TypeError("Dataflow retention requires a Session.")
        return purge_expired_runs(session, now=now, limit=limit)


__all__ = [
    "SqlDataflowRunWorker",
    "dispatch_pending_runs",
    "purge_expired_runs",
    "run_metrics",
]