from __future__ import annotations

from sqlalchemy import Index, Integer, String, Text, UniqueConstraint, text
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    pass


class HostRow(Base):
    __tablename__ = "host_registrations"
    __table_args__ = (
        UniqueConstraint(
            "agent_instance_id",
            name="uq_host_registrations_agent_instance_id",
        ),
        UniqueConstraint(
            "credential_digest",
            name="uq_host_registrations_credential_digest",
        ),
        UniqueConstraint(
            "enrollment_token_digest",
            name="uq_host_registrations_enrollment_token_digest",
        ),
    )

    host_id: Mapped[str] = mapped_column(String, primary_key=True)
    address: Mapped[str | None] = mapped_column(String, nullable=True)
    last_seen_at: Mapped[str] = mapped_column(String, nullable=False)
    agent_instance_id: Mapped[str | None] = mapped_column(
        String,
        nullable=True,
    )
    credential_digest: Mapped[str | None] = mapped_column(
        String,
        nullable=True,
    )
    enrollment_token_digest: Mapped[str | None] = mapped_column(
        String,
        nullable=True,
    )
    display_name: Mapped[str | None] = mapped_column(String, nullable=True)
    enrolled_at: Mapped[str | None] = mapped_column(String, nullable=True)
    revoked_at: Mapped[str | None] = mapped_column(String, nullable=True)


class DeviceEnrollmentRow(Base):
    __tablename__ = "device_enrollments"
    __table_args__ = (
        UniqueConstraint(
            "host_id",
            "local_device_id",
            name="uq_device_enrollments_host_local",
        ),
        Index("ix_device_enrollments_host_id", "host_id"),
    )

    device_id: Mapped[str] = mapped_column(String, primary_key=True)
    host_id: Mapped[str] = mapped_column(String, nullable=False)
    local_device_id: Mapped[str] = mapped_column(String, nullable=False)
    driver_type: Mapped[str] = mapped_column(String, nullable=False)
    name: Mapped[str | None] = mapped_column(String, nullable=True)
    capability_tags_json: Mapped[str] = mapped_column(Text, nullable=False)
    enrolled_at: Mapped[str] = mapped_column(String, nullable=False)
    revoked_at: Mapped[str | None] = mapped_column(String, nullable=True)


class PooledDeviceRow(Base):
    __tablename__ = "pooled_devices"

    host_id: Mapped[str] = mapped_column(String, primary_key=True)
    device_id: Mapped[str] = mapped_column(String, primary_key=True)
    driver_type: Mapped[str] = mapped_column(String, nullable=False)
    status: Mapped[str] = mapped_column(String, nullable=False)
    capability_tags_json: Mapped[str] = mapped_column(Text, nullable=False)
    synced_at: Mapped[str | None] = mapped_column(String, nullable=True)


class ScheduledTaskRow(Base):
    __tablename__ = "scheduled_tasks"

    id: Mapped[str] = mapped_column(String, primary_key=True)
    goal: Mapped[str | None] = mapped_column(Text, nullable=True)
    workflow_definition_id: Mapped[str | None] = mapped_column(String, nullable=True)
    constraints_json: Mapped[str] = mapped_column(Text, nullable=False)
    status: Mapped[str] = mapped_column(String, nullable=False)
    assigned_device_id: Mapped[str | None] = mapped_column(String, nullable=True)
    assigned_host_id: Mapped[str | None] = mapped_column(String, nullable=True)
    attempt_count: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        default=0,
        server_default=text("0"),
    )
    lease_id: Mapped[str | None] = mapped_column(String, nullable=True)
    lease_expires_at: Mapped[str | None] = mapped_column(String, nullable=True)
    failure_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
    result_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    updated_at: Mapped[str | None] = mapped_column(String, nullable=True)
    created_at: Mapped[str] = mapped_column(String, nullable=False)


class TaskAttemptRow(Base):
    __tablename__ = "task_attempts"

    task_id: Mapped[str] = mapped_column(String, primary_key=True)
    attempt: Mapped[int] = mapped_column(Integer, primary_key=True)
    lease_id: Mapped[str] = mapped_column(String, nullable=False)
    host_id: Mapped[str] = mapped_column(String, nullable=False)
    device_id: Mapped[str] = mapped_column(String, nullable=False)
    status: Mapped[str] = mapped_column(String, nullable=False)
    lease_expires_at: Mapped[str] = mapped_column(String, nullable=False)
    created_at: Mapped[str] = mapped_column(String, nullable=False)
    completed_at: Mapped[str | None] = mapped_column(String, nullable=True)
    failure_reason: Mapped[str | None] = mapped_column(Text, nullable=True)
    result_json: Mapped[str | None] = mapped_column(Text, nullable=True)


class PluginRow(Base):
    __tablename__ = "plugins"

    name: Mapped[str] = mapped_column(String, primary_key=True)
    version: Mapped[str] = mapped_column(String, nullable=False)
    entry_point_kind: Mapped[str] = mapped_column(String, nullable=False)
    target: Mapped[str] = mapped_column(String, nullable=False)
    wired: Mapped[int] = mapped_column(Integer, nullable=False)