from datetime import datetime
from uuid import UUID

from sqlalchemy import (
    JSON,
    Boolean,
    CheckConstraint,
    Computed,
    DateTime,
    ForeignKey,
    Index,
    Integer,
    PrimaryKeyConstraint,
    String,
    Text,
    func,
    text,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.dialects.postgresql import UUID as PG_UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    pass


def _jsonb():
    return JSONB().with_variant(JSON(), "sqlite")


def _uuid():
    return PG_UUID(as_uuid=True).with_variant(String(36), "sqlite")


class Agent(Base):
    __tablename__ = "agents"

    agent_id: Mapped[str] = mapped_column(Text, primary_key=True)
    hostname: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(Text, nullable=False, default="alive")
    last_seen_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    features: Mapped[dict] = mapped_column(
        _jsonb(),
        nullable=False,
        default=lambda: {"push": False, "metrics": False},
        server_default=text("jsonb_build_object('push', false, 'metrics', false)"),
    )
    labels: Mapped[dict] = mapped_column(
        _jsonb(), nullable=False, default=dict, server_default=text("'{}'::jsonb")
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    accepted_key_hash: Mapped[str | None] = mapped_column(Text, nullable=True)
    accepted_key_fingerprint: Mapped[str | None] = mapped_column(Text, nullable=True)
    accepted_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    pending_key_hash: Mapped[str | None] = mapped_column(Text, nullable=True)
    pending_key_fingerprint: Mapped[str | None] = mapped_column(Text, nullable=True)
    pending_hostname: Mapped[str | None] = mapped_column(Text, nullable=True)
    pending_seen_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    pending_features: Mapped[dict | None] = mapped_column(_jsonb(), nullable=True)
    pending_labels: Mapped[dict | None] = mapped_column(_jsonb(), nullable=True)
    admission_rank: Mapped[int] = mapped_column(
        Integer,
        Computed("CASE WHEN pending_key_hash IS NOT NULL THEN 0 ELSE 1 END", persisted=True),
        nullable=False,
    )

    __table_args__ = (
        CheckConstraint("status IN ('alive','stale','dead')", name="agents_status_chk"),
        Index("agents_status_idx", "status"),
        Index("agents_last_seen_idx", "last_seen_at"),
        Index("agents_admission_rank_idx", "admission_rank", "agent_id"),
        Index("agents_hostname_idx", "hostname", unique=True),
        Index(
            "agents_accepted_key_hash_idx",
            "accepted_key_hash",
            unique=True,
            postgresql_where=text("accepted_key_hash IS NOT NULL"),
        ),
        Index(
            "agents_pending_key_hash_idx",
            "pending_key_hash",
            unique=True,
            postgresql_where=text("pending_key_hash IS NOT NULL"),
        ),
    )


class AgentBlacklist(Base):
    __tablename__ = "agent_blacklist"

    key_hash: Mapped[str] = mapped_column(Text, primary_key=True)
    key_fingerprint: Mapped[str] = mapped_column(Text, nullable=False)
    agent_id: Mapped[str] = mapped_column(Text, nullable=False)
    hostname: Mapped[str] = mapped_column(Text, nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    last_seen_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)

    __table_args__ = (
        Index("agent_blacklist_fingerprint_idx", "key_fingerprint"),
        Index("agent_blacklist_last_seen_idx", "last_seen_at"),
    )


class Check(Base):
    __tablename__ = "checks"

    agent_id: Mapped[str] = mapped_column(
        Text, ForeignKey("agents.agent_id", ondelete="CASCADE"), nullable=False
    )
    check_id: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(Text, nullable=False)
    exit_code: Mapped[int] = mapped_column(Integer, nullable=False)
    last_observed_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    last_event_id: Mapped[UUID] = mapped_column(_uuid(), nullable=False)
    incident_key: Mapped[str] = mapped_column(Text, nullable=False)
    summary: Mapped[str | None] = mapped_column(Text, nullable=True)

    __table_args__ = (
        PrimaryKeyConstraint("agent_id", "check_id", name="checks_pkey"),
        CheckConstraint(
            "status IN ('ok','warning','critical','unknown')", name="checks_status_chk"
        ),
        Index("checks_status_idx", "status"),
    )


class Event(Base):
    __tablename__ = "events"

    event_id: Mapped[UUID] = mapped_column(_uuid(), nullable=False)
    agent_id: Mapped[str] = mapped_column(Text, nullable=False)
    check_id: Mapped[str] = mapped_column(Text, nullable=False)
    observed_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    # Partition key (RANGE) on Postgres; composite PK = (event_id, received_at).
    received_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    status: Mapped[str] = mapped_column(Text, nullable=False)
    exit_code: Mapped[int] = mapped_column(Integer, nullable=False)
    duration_ms: Mapped[int] = mapped_column(Integer, nullable=False)
    output: Mapped[str | None] = mapped_column(Text, nullable=True)
    output_truncated: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    notifications_enabled: Mapped[bool] = mapped_column(
        Boolean, nullable=False, default=True, server_default=text("TRUE")
    )
    incident_key: Mapped[str] = mapped_column(Text, nullable=False)

    __table_args__ = (
        PrimaryKeyConstraint("event_id", "received_at", name="events_pkey"),
        CheckConstraint(
            "status IN ('ok','warning','critical','unknown')", name="events_status_chk"
        ),
        CheckConstraint("duration_ms >= 0", name="events_duration_chk"),
        Index(
            "events_agent_check_observed_idx",
            "agent_id",
            "check_id",
            text("observed_at DESC"),
        ),
        Index("events_observed_at_idx", text("observed_at DESC")),
        Index("events_received_at_idx", text("received_at DESC")),
        Index("events_received_event_idx", text("received_at DESC"), text("event_id DESC")),
        Index(
            "events_agent_received_event_idx",
            "agent_id",
            text("received_at DESC"),
            text("event_id DESC"),
        ),
        Index(
            "events_check_received_event_idx",
            "check_id",
            text("received_at DESC"),
            text("event_id DESC"),
        ),
    )


class EventIngestDedup(Base):
    __tablename__ = "event_ingest_dedup"

    event_id: Mapped[UUID] = mapped_column(_uuid(), primary_key=True)
    received_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )

    __table_args__ = (Index("event_ingest_dedup_received_at_idx", "received_at"),)


class Incident(Base):
    __tablename__ = "incidents"

    id: Mapped[UUID] = mapped_column(_uuid(), primary_key=True)
    incident_key: Mapped[str] = mapped_column(Text, nullable=False)
    agent_id: Mapped[str] = mapped_column(Text, nullable=False)
    check_id: Mapped[str] = mapped_column(Text, nullable=False)
    state: Mapped[str] = mapped_column(Text, nullable=False)
    severity: Mapped[str] = mapped_column(Text, nullable=False)
    opened_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
    resolved_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    summary: Mapped[str | None] = mapped_column(Text, nullable=True)
    last_event_id: Mapped[UUID] = mapped_column(_uuid(), nullable=False)
    # PH-review: stored generated column so the default UI sort
    # (status_rank, opened_at, id) is served by a plain B-tree index without
    # planner mismatches between query-time and index-time CASE expressions.
    # Keep aligned with `_STATUS_RANK` in api/incidents.py.
    status_rank: Mapped[int] = mapped_column(
        Integer,
        Computed(
            "CASE "
            "WHEN state = 'open' AND severity = 'critical' THEN 5 "
            "WHEN state = 'open' AND severity = 'warning' THEN 4 "
            "WHEN state = 'open' AND severity = 'unknown' THEN 3 "
            "WHEN state = 'resolved' AND severity = 'critical' THEN 2 "
            "WHEN state = 'resolved' AND severity = 'warning' THEN 1 "
            "WHEN state = 'resolved' AND severity = 'unknown' THEN 0 "
            "END",
            persisted=True,
        ),
        nullable=False,
    )

    __table_args__ = (
        CheckConstraint("state IN ('open','resolved')", name="incidents_state_chk"),
        CheckConstraint(
            "severity IN ('warning','critical','unknown')", name="incidents_severity_chk"
        ),
        Index(
            "incidents_open_key_idx",
            "incident_key",
            unique=True,
            postgresql_where=text("state = 'open'"),
        ),
        Index(
            "incidents_status_rank_idx",
            text("status_rank DESC"),
            text("opened_at DESC"),
            text("id DESC"),
        ),
    )


class NotificationOutbox(Base):
    __tablename__ = "notification_outbox"

    id: Mapped[UUID] = mapped_column(_uuid(), primary_key=True)
    incident_id: Mapped[UUID] = mapped_column(
        _uuid(), ForeignKey("incidents.id", ondelete="CASCADE"), nullable=False
    )
    notifier: Mapped[str] = mapped_column(Text, nullable=False)
    event_type: Mapped[str] = mapped_column(Text, nullable=False)
    state: Mapped[str] = mapped_column(Text, nullable=False)
    attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    next_attempt_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    last_error: Mapped[str | None] = mapped_column(Text, nullable=True)
    payload: Mapped[dict] = mapped_column(_jsonb(), nullable=False)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now()
    )

    __table_args__ = (
        CheckConstraint("event_type IN ('firing','resolved')", name="outbox_event_type_chk"),
        CheckConstraint(
            "state IN ('pending','sending','sent','retry','failed','discarded')",
            name="outbox_state_chk",
        ),
        Index(
            "outbox_pending_idx",
            "state",
            "next_attempt_at",
            postgresql_where=text("state IN ('pending','retry')"),
        ),
        Index("outbox_created_at_idx", text("created_at DESC"), text("id DESC")),
        Index("outbox_updated_at_idx", text("updated_at DESC"), text("id DESC")),
    )