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 Session, sessionmaker

from cloud.db_models import (
    Base,
    HostRow,
    PluginRow,
    PooledDeviceRow,
    ScheduledTaskRow,
)
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,
                    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 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 = {}
    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,
        created_at=_parse_dt(row.created_at) or utc_now(),
    )


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),
    )