from __future__ import annotations

import json
from dataclasses import asdict
from datetime import datetime
from typing import Any

from sqlalchemy import Engine, delete, func, select
from sqlalchemy.orm import sessionmaker

from cloud.db_models import (
    Base,
    HostRow,
    PluginRow,
    PooledDeviceRow,
    ScheduledTaskRow,
    TaskAttemptRow,
)
from core.models import utc_now


class SQLAlchemyCloudRepository:
    """SQLAlchemy adapter preserving the existing CloudStore CRUD surface."""

    def __init__(self, engine: Engine, *, create_schema: bool = True) -> None:
        self.engine = engine
        self._sessions = sessionmaker(bind=engine, expire_on_commit=False)
        if create_schema:
            Base.metadata.create_all(engine)

    def upsert_host(
        self,
        host_id: str,
        *,
        address: str | None,
        last_seen_at: datetime,
    ) -> None:
        with self._sessions.begin() as session:
            row = session.get(HostRow, host_id)
            if row is None:
                session.add(
                    HostRow(
                        host_id=host_id,
                        address=address,
                        last_seen_at=_iso(last_seen_at),
                    )
                )
                return
            if address is not None:
                row.address = address
            row.last_seen_at = _iso(last_seen_at)

    def replace_host_devices(
        self,
        host_id: str,
        devices: list[Any],
    ) -> None:
        with self._sessions.begin() as session:
            session.execute(
                delete(PooledDeviceRow).where(PooledDeviceRow.host_id == host_id)
            )
            session.add_all(
                [
                    PooledDeviceRow(
                        device_id=device.device_id,
                        host_id=device.host_id,
                        driver_type=device.driver_type,
                        status=device.status,
                        capability_tags_json=json.dumps(
                            list(device.capability_tags),
                            ensure_ascii=False,
                        ),
                        synced_at=_iso(device.synced_at) if device.synced_at else None,
                    )
                    for device in devices
                ]
            )

    def list_hosts(self) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(select(HostRow).order_by(HostRow.host_id)).all()
        return [_host_from_row(row) for row in rows]

    def get_host(self, host_id: str) -> Any | None:
        with self._sessions() as session:
            row = session.get(HostRow, host_id)
        return _host_from_row(row) if row else None

    def list_devices(self) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(
                select(PooledDeviceRow).order_by(
                    PooledDeviceRow.host_id,
                    PooledDeviceRow.device_id,
                )
            ).all()
        return [_device_from_row(row) for row in rows]

    def get_device(self, device_id: str) -> Any | None:
        with self._sessions() as session:
            row = session.scalars(
                select(PooledDeviceRow)
                .where(PooledDeviceRow.device_id == device_id)
                .order_by(PooledDeviceRow.host_id)
                .limit(1)
            ).first()
        return _device_from_row(row) if row else None

    def enqueue_task(self, task: Any) -> None:
        with self._sessions.begin() as session:
            session.add(
                ScheduledTaskRow(
                    id=task.id,
                    goal=task.goal,
                    workflow_definition_id=task.workflow_definition_id,
                    constraints_json=json.dumps(
                        asdict(task.constraints),
                        ensure_ascii=False,
                    ),
                    status=task.status,
                    assigned_device_id=task.assigned_device_id,
                    assigned_host_id=task.assigned_host_id,
                    attempt_count=task.attempt_count,
                    lease_id=task.lease_id,
                    lease_expires_at=(
                        _iso(task.lease_expires_at) if task.lease_expires_at else None
                    ),
                    failure_reason=task.failure_reason,
                    result_json=(
                        json.dumps(task.terminal_result, ensure_ascii=False)
                        if task.terminal_result is not None
                        else None
                    ),
                    updated_at=_iso(task.updated_at) if task.updated_at else None,
                    created_at=_iso(task.created_at),
                )
            )

    def list_queued_tasks(self) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(
                select(ScheduledTaskRow)
                .where(ScheduledTaskRow.status == "queued")
                .order_by(ScheduledTaskRow.created_at, ScheduledTaskRow.id)
            ).all()
        return [_task_from_row(row) for row in rows]

    def get_task(self, task_id: str) -> Any | None:
        with self._sessions() as session:
            row = session.get(ScheduledTaskRow, task_id)
        return _task_from_row(row) if row else None

    def update_task(
        self,
        task_id: str,
        *,
        status: str | None = None,
        assigned_device_id: str | None = None,
        assigned_host_id: str | None = None,
    ) -> None:
        with self._sessions.begin() as session:
            row = session.get(ScheduledTaskRow, task_id)
            if row is None:
                return
            if status is not None:
                row.status = status
            if assigned_device_id is not None:
                row.assigned_device_id = assigned_device_id
            if assigned_host_id is not None:
                row.assigned_host_id = assigned_host_id

    def count_queued_tasks(self) -> int:
        with self._sessions() as session:
            count = session.scalar(
                select(func.count())
                .select_from(ScheduledTaskRow)
                .where(ScheduledTaskRow.status == "queued")
            )
        return int(count or 0)

    def save_plugin(self, manifest: Any, *, wired: bool) -> None:
        with self._sessions.begin() as session:
            row = session.get(PluginRow, manifest.name)
            if row is None:
                session.add(
                    PluginRow(
                        name=manifest.name,
                        version=manifest.version,
                        entry_point_kind=manifest.entry_point_kind,
                        target=manifest.target,
                        wired=1 if wired else 0,
                    )
                )
                return
            row.version = manifest.version
            row.entry_point_kind = manifest.entry_point_kind
            row.target = manifest.target
            row.wired = 1 if wired else 0

    def list_plugins(self) -> list[tuple[Any, bool]]:
        with self._sessions() as session:
            rows = session.scalars(select(PluginRow).order_by(PluginRow.name)).all()
        return [_plugin_from_row(row) for row in rows]

    def get_plugin(self, name: str) -> tuple[Any, bool] | None:
        with self._sessions() as session:
            row = session.get(PluginRow, name)
        return _plugin_from_row(row) if row else None

    def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
        with self._sessions() as session:
            device_ids = session.scalars(
                select(ScheduledTaskRow.assigned_device_id).where(
                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
                    ScheduledTaskRow.assigned_device_id.is_not(None),
                    ScheduledTaskRow.lease_expires_at.is_not(None),
                    ScheduledTaskRow.lease_expires_at > _iso(now),
                )
            ).all()
        return {device_id for device_id in device_ids if device_id is not None}

    def assign_task(
        self,
        *,
        task_id: str,
        host_id: str,
        device_id: str,
        lease_id: str,
        lease_expires_at: datetime,
        now: datetime,
    ) -> Any | None:
        with self._sessions.begin() as session:
            task_statement = select(ScheduledTaskRow).where(
                ScheduledTaskRow.id == task_id
            )
            device_statement = select(PooledDeviceRow).where(
                PooledDeviceRow.host_id == host_id,
                PooledDeviceRow.device_id == device_id,
            )
            if self.engine.dialect.name == "postgresql":
                task_statement = task_statement.with_for_update(skip_locked=True)
                device_statement = device_statement.with_for_update()

            task = session.scalars(task_statement).first()
            if task is None or task.status != "queued":
                return None
            device = session.scalars(device_statement).first()
            if device is None or device.status != "idle":
                return None

            active_reservation = session.scalar(
                select(ScheduledTaskRow.id)
                .where(
                    ScheduledTaskRow.id != task_id,
                    ScheduledTaskRow.assigned_device_id == device_id,
                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
                    ScheduledTaskRow.lease_expires_at.is_not(None),
                    ScheduledTaskRow.lease_expires_at > _iso(now),
                )
                .limit(1)
            )
            if active_reservation is not None:
                return None

            attempt = task.attempt_count + 1
            task.status = "assigned"
            task.assigned_host_id = host_id
            task.assigned_device_id = device_id
            task.attempt_count = attempt
            task.lease_id = lease_id
            task.lease_expires_at = _iso(lease_expires_at)
            task.failure_reason = None
            task.result_json = None
            task.updated_at = _iso(now)
            session.add(
                TaskAttemptRow(
                    task_id=task.id,
                    attempt=attempt,
                    lease_id=lease_id,
                    host_id=host_id,
                    device_id=device_id,
                    status="assigned",
                    lease_expires_at=_iso(lease_expires_at),
                    created_at=_iso(now),
                    completed_at=None,
                    failure_reason=None,
                    result_json=None,
                )
            )
            session.flush()
            return _leased_assignment_from_row(task)

    def claim_assignment(
        self,
        *,
        host_id: str,
        now: datetime,
    ) -> Any | None:
        with self._sessions.begin() as session:
            statement = (
                select(ScheduledTaskRow)
                .where(
                    ScheduledTaskRow.status == "assigned",
                    ScheduledTaskRow.assigned_host_id == host_id,
                    ScheduledTaskRow.lease_id.is_not(None),
                    ScheduledTaskRow.lease_expires_at.is_not(None),
                    ScheduledTaskRow.lease_expires_at > _iso(now),
                )
                .order_by(ScheduledTaskRow.created_at, ScheduledTaskRow.id)
                .limit(1)
            )
            if self.engine.dialect.name == "postgresql":
                statement = statement.with_for_update(skip_locked=True)

            task = session.scalars(statement).first()
            if task is None:
                return None
            attempt = session.get(
                TaskAttemptRow,
                (task.id, task.attempt_count),
                with_for_update=self.engine.dialect.name == "postgresql",
            )
            if (
                attempt is None
                or attempt.status != "assigned"
                or attempt.lease_id != task.lease_id
            ):
                return None

            task.status = "dispatched"
            task.updated_at = _iso(now)
            attempt.status = "dispatched"
            session.flush()
            return _leased_assignment_from_row(task)

    def renew_lease(
        self,
        *,
        task_id: str,
        attempt: int,
        lease_id: str,
        host_id: str,
        lease_expires_at: datetime,
        now: datetime,
    ) -> str:
        with self._sessions.begin() as session:
            task = session.get(
                ScheduledTaskRow,
                task_id,
                with_for_update=self.engine.dialect.name == "postgresql",
            )
            if task is None:
                return "not_found"
            if (
                task.status not in {"assigned", "dispatched"}
                or task.attempt_count != attempt
                or task.lease_id != lease_id
                or task.assigned_host_id != host_id
            ):
                return "conflict"
            current_expiry = _parse_dt(task.lease_expires_at)
            if current_expiry is None or current_expiry <= now:
                return "expired"
            if lease_expires_at <= now:
                return "conflict"

            attempt_row = session.get(
                TaskAttemptRow,
                (task_id, attempt),
                with_for_update=self.engine.dialect.name == "postgresql",
            )
            if (
                attempt_row is None
                or attempt_row.status not in {"assigned", "dispatched"}
                or attempt_row.lease_id != lease_id
                or attempt_row.host_id != host_id
            ):
                return "conflict"

            renewed_until = _iso(lease_expires_at)
            task.lease_expires_at = renewed_until
            task.updated_at = _iso(now)
            attempt_row.lease_expires_at = renewed_until
            return "renewed"

    def record_task_result(
        self,
        *,
        task_id: str,
        attempt: int,
        lease_id: str,
        host_id: str,
        status: str,
        failure_reason: str | None,
        terminal_result: dict[str, Any] | None,
        completed_at: datetime,
    ) -> str:
        with self._sessions.begin() as session:
            task = session.get(
                ScheduledTaskRow,
                task_id,
                with_for_update=self.engine.dialect.name == "postgresql",
            )
            if task is None:
                return "conflict"
            attempt_row = session.get(
                TaskAttemptRow,
                (task_id, attempt),
                with_for_update=self.engine.dialect.name == "postgresql",
            )
            if (
                attempt_row is None
                or task.attempt_count != attempt
                or task.lease_id != lease_id
                or task.assigned_host_id != host_id
                or attempt_row.lease_id != lease_id
                or attempt_row.host_id != host_id
            ):
                return "conflict"

            if task.status in {"done", "failed"}:
                if (
                    task.status == status
                    and task.failure_reason == failure_reason
                    and _parse_json_object(task.result_json) == terminal_result
                    and attempt_row.status == status
                ):
                    return "already_recorded"
                return "conflict"
            if task.status not in {"assigned", "dispatched"}:
                return "conflict"
            current_expiry = _parse_dt(task.lease_expires_at)
            if current_expiry is None or current_expiry <= completed_at:
                return "conflict"
            if status not in {"done", "failed"}:
                return "conflict"

            result_json = (
                json.dumps(terminal_result, ensure_ascii=False)
                if terminal_result is not None
                else None
            )
            completed_at_iso = _iso(completed_at)
            task.status = status
            task.failure_reason = failure_reason
            task.result_json = result_json
            task.updated_at = completed_at_iso
            attempt_row.status = status
            attempt_row.completed_at = completed_at_iso
            attempt_row.failure_reason = failure_reason
            attempt_row.result_json = result_json
            return "recorded"

    def list_task_attempts(self, task_id: str) -> list[Any]:
        with self._sessions() as session:
            rows = session.scalars(
                select(TaskAttemptRow)
                .where(TaskAttemptRow.task_id == task_id)
                .order_by(TaskAttemptRow.attempt)
            ).all()
        return [_task_attempt_from_row(row) for row in rows]

    def health_check(self) -> None:
        with self._sessions() as session:
            session.execute(select(1))

    def close(self) -> None:
        self.engine.dispose()


def _iso(value: datetime) -> str:
    return value.isoformat()


def _parse_dt(value: str | None) -> datetime | None:
    if not value:
        return None
    try:
        return datetime.fromisoformat(value)
    except ValueError:
        return None


def _host_from_row(row: HostRow) -> Any:
    from cloud.pool import HostRegistration

    return HostRegistration(
        host_id=row.host_id,
        address=row.address,
        last_seen_at=_parse_dt(row.last_seen_at) or utc_now(),
    )


def _device_from_row(row: PooledDeviceRow) -> Any:
    from cloud.pool import PooledDevice

    try:
        tags = list(json.loads(row.capability_tags_json))
    except TypeError, ValueError:
        tags = []
    return PooledDevice(
        device_id=row.device_id,
        host_id=row.host_id,
        driver_type=row.driver_type,
        status=row.status,
        capability_tags=tags,
        synced_at=_parse_dt(row.synced_at),
    )


def _task_from_row(row: ScheduledTaskRow) -> Any:
    from cloud.scheduler import ScheduledTask, TaskConstraints

    try:
        constraints_data = json.loads(row.constraints_json)
    except TypeError, ValueError:
        constraints_data = {}
    terminal_result = _parse_json_object(row.result_json)
    return ScheduledTask(
        id=row.id,
        goal=row.goal,
        workflow_definition_id=row.workflow_definition_id,
        constraints=TaskConstraints(
            driver_type=constraints_data.get("driver_type"),
            capability_tags=list(constraints_data.get("capability_tags") or []),
        ),
        status=row.status,
        assigned_device_id=row.assigned_device_id,
        assigned_host_id=row.assigned_host_id,
        attempt_count=row.attempt_count,
        lease_id=row.lease_id,
        lease_expires_at=_parse_dt(row.lease_expires_at),
        terminal_result=terminal_result,
        failure_reason=row.failure_reason,
        updated_at=_parse_dt(row.updated_at),
        created_at=_parse_dt(row.created_at) or utc_now(),
    )


def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
    from cloud.repository import TaskAttemptRecord

    return TaskAttemptRecord(
        task_id=row.task_id,
        attempt=row.attempt,
        lease_id=row.lease_id,
        host_id=row.host_id,
        device_id=row.device_id,
        status=row.status,
        lease_expires_at=_parse_dt(row.lease_expires_at) or utc_now(),
        created_at=_parse_dt(row.created_at) or utc_now(),
        completed_at=_parse_dt(row.completed_at),
        failure_reason=row.failure_reason,
        terminal_result=_parse_json_object(row.result_json),
    )


def _leased_assignment_from_row(row: ScheduledTaskRow) -> Any:
    from cloud.repository import LeasedAssignment

    lease_expires_at = _parse_dt(row.lease_expires_at)
    if (
        row.lease_id is None
        or lease_expires_at is None
        or row.assigned_host_id is None
        or row.assigned_device_id is None
    ):
        raise ValueError(f"task {row.id!r} does not contain a complete lease")
    return LeasedAssignment(
        task_id=row.id,
        attempt=row.attempt_count,
        lease_id=row.lease_id,
        lease_expires_at=lease_expires_at,
        host_id=row.assigned_host_id,
        device_id=row.assigned_device_id,
        goal=row.goal,
        workflow_definition_id=row.workflow_definition_id,
    )


def _parse_json_object(value: str | None) -> dict[str, Any] | None:
    if value is None:
        return None
    try:
        parsed = json.loads(value)
    except TypeError, ValueError:
        return None
    return parsed if isinstance(parsed, dict) else None


def _plugin_from_row(row: PluginRow) -> tuple[Any, bool]:
    from cloud.plugins import PluginManifest

    return (
        PluginManifest(
            name=row.name,
            version=row.version,
            entry_point_kind=row.entry_point_kind,
            target=row.target,
        ),
        bool(row.wired),
    )