from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from enum import StrEnum
import hashlib
import json
import os
import socket
from typing import Any
from uuid import uuid4

from sqlalchemy import BigInteger, DateTime, Index, Integer, JSON, String, UniqueConstraint, select
from sqlalchemy.orm import Mapped, Session, mapped_column

from govoplan_core.db.base import Base, TimestampMixin, utcnow


class RuntimeNodeState(StrEnum):
    ACTIVE = "active"
    DRAINING = "draining"
    STOPPED = "stopped"


class RuntimeCoordinationError(RuntimeError):
    pass


class LeaseUnavailable(RuntimeCoordinationError):
    pass


class StaleFence(RuntimeCoordinationError):
    pass


class RuntimeNode(Base, TimestampMixin):
    __tablename__ = "core_runtime_nodes"
    __table_args__ = (
        UniqueConstraint(
            "installation_id",
            "node_id",
            name="uq_core_runtime_node_installation_node",
        ),
        Index(
            "ix_core_runtime_nodes_installation_state_heartbeat",
            "installation_id",
            "state",
            "last_heartbeat_at",
        ),
    )

    id: Mapped[str] = mapped_column(
        String(36),
        primary_key=True,
        default=lambda: str(uuid4()),
    )
    installation_id: Mapped[str] = mapped_column(String(100), nullable=False, index=True)
    node_id: Mapped[str] = mapped_column(String(200), nullable=False, index=True)
    incarnation: Mapped[str] = mapped_column(String(36), nullable=False, index=True)
    role: Mapped[str] = mapped_column(String(40), nullable=False, index=True)
    software_version: Mapped[str] = mapped_column(String(80), nullable=False)
    composition_hash: Mapped[str] = mapped_column(String(64), nullable=False, index=True)
    queues: Mapped[list[str]] = mapped_column(JSON, default=list, nullable=False)
    state: Mapped[str] = mapped_column(
        String(30),
        default=RuntimeNodeState.ACTIVE.value,
        nullable=False,
        index=True,
    )
    started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow, nullable=False)
    last_heartbeat_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow, nullable=False)
    drain_requested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    drain_reason: Mapped[str | None] = mapped_column(String(500))
    stopped_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict, nullable=False)


class DistributedLease(Base, TimestampMixin):
    __tablename__ = "core_distributed_leases"
    __table_args__ = (
        UniqueConstraint(
            "installation_id",
            "resource_key",
            name="uq_core_distributed_lease_resource",
        ),
        Index(
            "ix_core_distributed_leases_expiry",
            "installation_id",
            "expires_at",
        ),
    )

    id: Mapped[str] = mapped_column(
        String(36),
        primary_key=True,
        default=lambda: str(uuid4()),
    )
    installation_id: Mapped[str] = mapped_column(String(100), nullable=False, index=True)
    resource_key: Mapped[str] = mapped_column(String(255), nullable=False, index=True)
    holder_node_id: Mapped[str | None] = mapped_column(String(200), index=True)
    holder_incarnation: Mapped[str | None] = mapped_column(String(36), index=True)
    fencing_token: Mapped[int] = mapped_column(
        BigInteger().with_variant(Integer, "sqlite"),
        default=0,
        nullable=False,
    )
    acquired_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    renewed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow, nullable=False)
    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata", JSON, default=dict, nullable=False)


@dataclass(frozen=True, slots=True)
class RuntimeIdentity:
    installation_id: str
    node_id: str
    incarnation: str
    role: str
    software_version: str
    composition_hash: str
    queues: tuple[str, ...] = ()


@dataclass(frozen=True, slots=True)
class LeaseClaim:
    installation_id: str
    resource_key: str
    holder_node_id: str
    holder_incarnation: str
    fencing_token: int
    expires_at: datetime


def runtime_identity(
    settings: object,
    *,
    software_version: str,
    module_ids: tuple[str, ...] = (),
    role: str | None = None,
    node_id: str | None = None,
    incarnation: str | None = None,
    queues: tuple[str, ...] | None = None,
) -> RuntimeIdentity:
    effective_role = str(
        role or getattr(settings, "runtime_role", "api") or "api"
    ).strip().lower()
    effective_node_id = str(
        node_id
        or getattr(settings, "runtime_node_id", None)
        or os.getenv("HOSTNAME")
        or socket.gethostname()
    ).strip()
    installation_id = str(
        getattr(settings, "installation_id", "govoplan-local")
        or "govoplan-local"
    ).strip()
    if not installation_id or not effective_node_id or not effective_role:
        raise RuntimeCoordinationError(
            "Runtime identity requires installation, node, and role identifiers"
        )
    effective_queues = queues
    if effective_queues is None:
        effective_queues = tuple(
            item.strip()
            for item in str(getattr(settings, "celery_queues", "") or "").split(",")
            if item.strip()
        )
    composition_hash = hashlib.sha256(
        json.dumps(sorted(module_ids), separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    return RuntimeIdentity(
        installation_id=installation_id,
        node_id=effective_node_id,
        incarnation=incarnation or str(uuid4()),
        role=effective_role,
        software_version=str(software_version),
        composition_hash=composition_hash,
        queues=tuple(sorted(set(effective_queues))),
    )


def register_runtime_node(
    session: Session,
    identity: RuntimeIdentity,
    *,
    metadata: dict[str, Any] | None = None,
    now: datetime | None = None,
) -> RuntimeNode:
    observed_at = _as_utc(now or utcnow())
    _insert_node_placeholder(session, identity, observed_at)
    node = _locked_node(session, identity.installation_id, identity.node_id)
    if node is None:  # pragma: no cover - guarded by insert/upsert
        raise RuntimeCoordinationError("Runtime node registration was not persisted")
    node.incarnation = identity.incarnation
    node.role = identity.role
    node.software_version = identity.software_version
    node.composition_hash = identity.composition_hash
    node.queues = list(identity.queues)
    node.state = RuntimeNodeState.ACTIVE.value
    node.started_at = observed_at
    node.last_heartbeat_at = observed_at
    node.drain_requested_at = None
    node.drain_reason = None
    node.stopped_at = None
    node.metadata_ = dict(metadata or {})
    session.add(node)
    session.flush()
    return node


def heartbeat_runtime_node(
    session: Session,
    identity: RuntimeIdentity,
    *,
    now: datetime | None = None,
    metadata: dict[str, Any] | None = None,
) -> RuntimeNode:
    node = _locked_node(session, identity.installation_id, identity.node_id)
    if node is None or node.incarnation != identity.incarnation:
        raise RuntimeCoordinationError(
            "Runtime heartbeat was rejected for a stale node incarnation"
        )
    if node.state == RuntimeNodeState.STOPPED.value:
        raise RuntimeCoordinationError("Stopped runtime node cannot heartbeat")
    node.last_heartbeat_at = _as_utc(now or utcnow())
    if metadata is not None:
        node.metadata_ = dict(metadata)
    session.add(node)
    session.flush()
    return node


def request_runtime_node_drain(
    session: Session,
    *,
    installation_id: str,
    node_id: str,
    reason: str | None = None,
    now: datetime | None = None,
) -> RuntimeNode:
    node = _locked_node(session, installation_id, node_id)
    if node is None:
        raise RuntimeCoordinationError("Runtime node was not found")
    if node.state == RuntimeNodeState.STOPPED.value:
        raise RuntimeCoordinationError("Stopped runtime node cannot be drained")
    node.state = RuntimeNodeState.DRAINING.value
    node.drain_requested_at = _as_utc(now or utcnow())
    node.drain_reason = str(reason or "operator request")[:500]
    session.add(node)
    session.flush()
    return node


def cancel_runtime_node_drain(
    session: Session,
    *,
    installation_id: str,
    node_id: str,
) -> RuntimeNode:
    node = _locked_node(session, installation_id, node_id)
    if node is None:
        raise RuntimeCoordinationError("Runtime node was not found")
    if node.state != RuntimeNodeState.DRAINING.value:
        return node
    node.state = RuntimeNodeState.ACTIVE.value
    node.drain_requested_at = None
    node.drain_reason = None
    session.add(node)
    session.flush()
    return node


def stop_runtime_node(
    session: Session,
    identity: RuntimeIdentity,
    *,
    now: datetime | None = None,
) -> bool:
    node = _locked_node(session, identity.installation_id, identity.node_id)
    if node is None or node.incarnation != identity.incarnation:
        return False
    observed_at = _as_utc(now or utcnow())
    node.state = RuntimeNodeState.STOPPED.value
    node.stopped_at = observed_at
    node.last_heartbeat_at = observed_at
    session.add(node)
    session.flush()
    return True


def runtime_node_is_draining(
    session: Session,
    identity: RuntimeIdentity,
) -> bool:
    node = session.execute(
        select(RuntimeNode).where(
            RuntimeNode.installation_id == identity.installation_id,
            RuntimeNode.node_id == identity.node_id,
            RuntimeNode.incarnation == identity.incarnation,
        )
    ).scalar_one_or_none()
    return node is not None and node.state == RuntimeNodeState.DRAINING.value


def list_runtime_nodes(
    session: Session,
    *,
    installation_id: str,
    stale_after_seconds: int,
    now: datetime | None = None,
) -> list[dict[str, Any]]:
    observed_at = _as_utc(now or utcnow())
    stale_before = observed_at - timedelta(seconds=max(1, stale_after_seconds))
    nodes = session.execute(
        select(RuntimeNode)
        .where(RuntimeNode.installation_id == installation_id)
        .order_by(RuntimeNode.role, RuntimeNode.node_id)
    ).scalars().all()
    return [
        {
            "node_id": node.node_id,
            "incarnation": node.incarnation,
            "role": node.role,
            "software_version": node.software_version,
            "composition_hash": node.composition_hash,
            "queues": list(node.queues or []),
            "state": node.state,
            "started_at": _iso(node.started_at),
            "last_heartbeat_at": _iso(node.last_heartbeat_at),
            "drain_requested_at": _iso(node.drain_requested_at),
            "drain_reason": node.drain_reason,
            "stopped_at": _iso(node.stopped_at),
            "stale": (
                node.state != RuntimeNodeState.STOPPED.value
                and _as_utc(node.last_heartbeat_at) < stale_before
            ),
            "metadata": dict(node.metadata_ or {}),
        }
        for node in nodes
    ]


def acquire_lease(
    session: Session,
    *,
    installation_id: str,
    resource_key: str,
    holder_node_id: str,
    holder_incarnation: str,
    ttl_seconds: int,
    metadata: dict[str, Any] | None = None,
    now: datetime | None = None,
) -> LeaseClaim | None:
    if ttl_seconds < 1:
        raise ValueError("Lease TTL must be at least one second")
    observed_at = _as_utc(now or utcnow())
    _insert_lease_placeholder(session, installation_id, resource_key, observed_at)
    lease = _locked_lease(session, installation_id, resource_key)
    if lease is None:  # pragma: no cover - guarded by insert/upsert
        raise RuntimeCoordinationError("Distributed lease was not persisted")
    same_holder = (
        lease.holder_node_id == holder_node_id
        and lease.holder_incarnation == holder_incarnation
    )
    expired = _as_utc(lease.expires_at) <= observed_at
    if lease.holder_node_id is not None and not same_holder and not expired:
        return None
    if not same_holder or expired:
        lease.fencing_token = int(lease.fencing_token or 0) + 1
        lease.acquired_at = observed_at
    lease.holder_node_id = holder_node_id
    lease.holder_incarnation = holder_incarnation
    lease.renewed_at = observed_at
    lease.expires_at = observed_at + timedelta(seconds=ttl_seconds)
    lease.metadata_ = dict(metadata or {})
    session.add(lease)
    session.flush()
    return _lease_claim(lease)


def renew_lease(
    session: Session,
    claim: LeaseClaim,
    *,
    ttl_seconds: int,
    now: datetime | None = None,
) -> LeaseClaim:
    if ttl_seconds < 1:
        raise ValueError("Lease TTL must be at least one second")
    observed_at = _as_utc(now or utcnow())
    lease = _locked_lease(session, claim.installation_id, claim.resource_key)
    _assert_matching_fence(lease, claim, observed_at=observed_at)
    lease.renewed_at = observed_at
    lease.expires_at = observed_at + timedelta(seconds=ttl_seconds)
    session.add(lease)
    session.flush()
    return _lease_claim(lease)


def release_lease(
    session: Session,
    claim: LeaseClaim,
    *,
    now: datetime | None = None,
) -> None:
    observed_at = _as_utc(now or utcnow())
    lease = _locked_lease(session, claim.installation_id, claim.resource_key)
    _assert_matching_fence(lease, claim, observed_at=observed_at)
    lease.holder_node_id = None
    lease.holder_incarnation = None
    lease.renewed_at = observed_at
    lease.expires_at = observed_at
    session.add(lease)
    session.flush()


def assert_lease_fence(
    session: Session,
    claim: LeaseClaim,
    *,
    now: datetime | None = None,
) -> None:
    lease = _locked_lease(session, claim.installation_id, claim.resource_key)
    _assert_matching_fence(
        lease,
        claim,
        observed_at=_as_utc(now or utcnow()),
    )


def _insert_node_placeholder(
    session: Session,
    identity: RuntimeIdentity,
    observed_at: datetime,
) -> None:
    values = {
        "id": str(uuid4()),
        "installation_id": identity.installation_id,
        "node_id": identity.node_id,
        "incarnation": identity.incarnation,
        "role": identity.role,
        "software_version": identity.software_version,
        "composition_hash": identity.composition_hash,
        "queues": list(identity.queues),
        "state": RuntimeNodeState.ACTIVE.value,
        "started_at": observed_at,
        "last_heartbeat_at": observed_at,
        "metadata": {},
        "created_at": observed_at,
        "updated_at": observed_at,
    }
    _insert_do_nothing(
        session,
        RuntimeNode.__table__,
        values,
        conflict_columns=("installation_id", "node_id"),
    )


def _insert_lease_placeholder(
    session: Session,
    installation_id: str,
    resource_key: str,
    observed_at: datetime,
) -> None:
    values = {
        "id": str(uuid4()),
        "installation_id": installation_id,
        "resource_key": resource_key,
        "fencing_token": 0,
        "expires_at": observed_at,
        "metadata": {},
        "created_at": observed_at,
        "updated_at": observed_at,
    }
    _insert_do_nothing(
        session,
        DistributedLease.__table__,
        values,
        conflict_columns=("installation_id", "resource_key"),
    )


def _insert_do_nothing(
    session: Session,
    table: Any,
    values: dict[str, Any],
    *,
    conflict_columns: tuple[str, ...],
) -> None:
    dialect = session.get_bind().dialect.name
    if dialect == "postgresql":
        from sqlalchemy.dialects.postgresql import insert

        statement = insert(table).values(**values).on_conflict_do_nothing(
            index_elements=list(conflict_columns)
        )
    elif dialect == "sqlite":
        from sqlalchemy.dialects.sqlite import insert

        statement = insert(table).values(**values).on_conflict_do_nothing(
            index_elements=list(conflict_columns)
        )
    else:  # pragma: no cover - GovOPlaN supports PostgreSQL and dev SQLite
        raise RuntimeCoordinationError(
            f"Distributed coordination is unsupported on {dialect!r}"
        )
    session.execute(statement)
    session.flush()


def _locked_node(
    session: Session,
    installation_id: str,
    node_id: str,
) -> RuntimeNode | None:
    return session.execute(
        select(RuntimeNode)
        .where(
            RuntimeNode.installation_id == installation_id,
            RuntimeNode.node_id == node_id,
        )
        .with_for_update()
    ).scalar_one_or_none()


def _locked_lease(
    session: Session,
    installation_id: str,
    resource_key: str,
) -> DistributedLease | None:
    return session.execute(
        select(DistributedLease)
        .where(
            DistributedLease.installation_id == installation_id,
            DistributedLease.resource_key == resource_key,
        )
        .with_for_update()
    ).scalar_one_or_none()


def _assert_matching_fence(
    lease: DistributedLease | None,
    claim: LeaseClaim,
    *,
    observed_at: datetime,
) -> None:
    if (
        lease is None
        or lease.holder_node_id != claim.holder_node_id
        or lease.holder_incarnation != claim.holder_incarnation
        or int(lease.fencing_token) != claim.fencing_token
        or _as_utc(lease.expires_at) <= observed_at
    ):
        raise StaleFence(
            f"Lease fence is stale for resource {claim.resource_key!r}"
        )


def _lease_claim(lease: DistributedLease) -> LeaseClaim:
    if lease.holder_node_id is None or lease.holder_incarnation is None:
        raise LeaseUnavailable("Distributed lease has no active holder")
    return LeaseClaim(
        installation_id=lease.installation_id,
        resource_key=lease.resource_key,
        holder_node_id=lease.holder_node_id,
        holder_incarnation=lease.holder_incarnation,
        fencing_token=int(lease.fencing_token),
        expires_at=_as_utc(lease.expires_at),
    )


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


def _iso(value: datetime | None) -> str | None:
    return _as_utc(value).isoformat() if value is not None else None


__all__ = [
    "DistributedLease",
    "LeaseClaim",
    "LeaseUnavailable",
    "RuntimeCoordinationError",
    "RuntimeIdentity",
    "RuntimeNode",
    "RuntimeNodeState",
    "StaleFence",
    "acquire_lease",
    "assert_lease_fence",
    "cancel_runtime_node_drain",
    "heartbeat_runtime_node",
    "list_runtime_nodes",
    "register_runtime_node",
    "release_lease",
    "renew_lease",
    "request_runtime_node_drain",
    "runtime_identity",
    "runtime_node_is_draining",
    "stop_runtime_node",
]
