from __future__ import annotations

import asyncio
import logging
import sqlite3
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import aiosqlite

logger = logging.getLogger(__name__)

ACTIVE_STATUSES = ("queued", "running")


def utc_now() -> str:
    return datetime.now(UTC).isoformat()


class UserAlreadyQueued(RuntimeError):
    def __init__(self, user_id: int) -> None:
        super().__init__(f"У пользователя {user_id} уже есть активный запрос")
        self.user_id = user_id


@dataclass(frozen=True)
class Request:
    id: int
    user_id: int
    timestamp: str
    prompt_text: str
    input_file_path: str | None
    output_file_path: str | None
    status: str
    comfy_prompt_id: str | None
    mode: str
    seed: int | None
    error_text: str | None
    finished_at: str | None

    @classmethod
    def from_row(cls, row: aiosqlite.Row) -> Request:
        return cls(
            id=row["id"],
            user_id=row["user_id"],
            timestamp=row["timestamp"],
            prompt_text=row["prompt_text"],
            input_file_path=row["input_file_path"],
            output_file_path=row["output_file_path"],
            status=row["status"],
            comfy_prompt_id=row["comfy_prompt_id"],
            mode=row["mode"],
            seed=row["seed"],
            error_text=row["error_text"],
            finished_at=row["finished_at"],
        )


SCHEMA = """
CREATE TABLE IF NOT EXISTS requests (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id          INTEGER NOT NULL,
    timestamp        TEXT    NOT NULL,
    prompt_text      TEXT    NOT NULL,
    input_file_path  TEXT,
    output_file_path TEXT,
    status           TEXT    NOT NULL DEFAULT 'queued'
                     CHECK (status IN ('queued', 'running', 'done', 'failed')),
    comfy_prompt_id  TEXT,
    mode             TEXT    NOT NULL CHECK (mode IN ('t2i', 'i2i')),
    seed             INTEGER,
    error_text       TEXT,
    finished_at      TEXT
);

CREATE INDEX IF NOT EXISTS idx_requests_user_status
    ON requests(user_id, status);

CREATE INDEX IF NOT EXISTS idx_requests_status_id
    ON requests(status, id);

CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active
    ON requests(user_id) WHERE status IN ('queued', 'running');
"""


class Database:
    def __init__(self, path: Path) -> None:
        self._path = path
        self._conn: aiosqlite.Connection | None = None
        self._write_lock = asyncio.Lock()

    @property
    def path(self) -> Path:
        return self._path

    async def connect(self) -> None:
        self._path.parent.mkdir(parents=True, exist_ok=True)
        self._conn = await aiosqlite.connect(self._path)
        self._conn.row_factory = aiosqlite.Row
        await self._conn.execute("PRAGMA journal_mode=WAL")
        await self._conn.execute("PRAGMA busy_timeout=5000")
        await self._conn.execute("PRAGMA foreign_keys=ON")
        await self._conn.executescript(SCHEMA)
        await self._conn.commit()

    async def close(self) -> None:
        if self._conn is not None:
            await self._conn.close()
            self._conn = None

    def _connection(self) -> aiosqlite.Connection:
        if self._conn is None:
            raise RuntimeError("База данных не подключена")
        return self._conn

    async def create_request(
        self,
        *,
        user_id: int,
        mode: str,
        prompt_text: str,
        input_file_path: str | None,
        seed: int,
    ) -> Request:
        conn = self._connection()
        async with self._write_lock:
            try:
                cursor = await conn.execute(
                    """
                    INSERT INTO requests
                        (user_id, timestamp, prompt_text, input_file_path, status, mode, seed)
                    VALUES (?, ?, ?, ?, 'queued', ?, ?)
                    """,
                    (user_id, utc_now(), prompt_text, input_file_path, mode, seed),
                )
                request_id = cursor.lastrowid
                await conn.commit()
            except sqlite3.IntegrityError as exc:
                await conn.rollback()
                if "UNIQUE" in str(exc).upper():
                    raise UserAlreadyQueued(user_id) from exc
                raise
        request = await self.get_request(int(request_id))
        if request is None:
            raise RuntimeError("Не удалось прочитать созданный запрос")
        return request

    async def get_request(self, request_id: int) -> Request | None:
        conn = self._connection()
        async with conn.execute("SELECT * FROM requests WHERE id = ?", (request_id,)) as cursor:
            row = await cursor.fetchone()
        return Request.from_row(row) if row is not None else None

    async def get_active_by_user(self, user_id: int) -> Request | None:
        conn = self._connection()
        async with conn.execute(
            """
            SELECT * FROM requests
            WHERE user_id = ? AND status IN ('queued', 'running')
            ORDER BY id DESC LIMIT 1
            """,
            (user_id,),
        ) as cursor:
            row = await cursor.fetchone()
        return Request.from_row(row) if row is not None else None

    async def list_queued(self) -> list[Request]:
        conn = self._connection()
        async with conn.execute(
            "SELECT * FROM requests WHERE status = 'queued' ORDER BY id"
        ) as cursor:
            rows = await cursor.fetchall()
        return [Request.from_row(row) for row in rows]

    async def count_ahead(self, request_id: int) -> int:
        conn = self._connection()
        async with conn.execute(
            """
            SELECT COUNT(*) AS total FROM requests
            WHERE id < ? AND status IN ('queued', 'running')
            """,
            (request_id,),
        ) as cursor:
            row = await cursor.fetchone()
        return int(row["total"]) if row is not None else 0

    async def set_running(self, request_id: int, comfy_prompt_id: str | None = None) -> None:
        await self._update(
            """
            UPDATE requests
            SET status = 'running',
                comfy_prompt_id = COALESCE(?, comfy_prompt_id),
                error_text = NULL
            WHERE id = ?
            """,
            (comfy_prompt_id, request_id),
        )

    async def set_comfy_prompt_id(self, request_id: int, comfy_prompt_id: str) -> None:
        await self._update(
            "UPDATE requests SET comfy_prompt_id = ? WHERE id = ?",
            (comfy_prompt_id, request_id),
        )

    async def set_done(self, request_id: int, output_file_path: str) -> None:
        await self._update(
            """
            UPDATE requests
            SET status = 'done',
                output_file_path = ?,
                error_text = NULL,
                finished_at = ?
            WHERE id = ?
            """,
            (output_file_path, utc_now(), request_id),
        )

    async def set_failed(self, request_id: int, error_text: str) -> None:
        await self._update(
            """
            UPDATE requests
            SET status = 'failed',
                error_text = ?,
                finished_at = ?
            WHERE id = ?
            """,
            (error_text[:1000], utc_now(), request_id),
        )

    async def recover(self, reason: str = "прервано перезапуском") -> list[Request]:
        """Помечает running как failed и возвращает queued для повторной постановки."""
        conn = self._connection()
        async with self._write_lock:
            await conn.execute(
                """
                UPDATE requests
                SET status = 'failed', error_text = ?, finished_at = ?
                WHERE status = 'running'
                """,
                (reason, utc_now()),
            )
            await conn.commit()
        return await self.list_queued()

    async def _update(self, sql: str, params: tuple[Any, ...]) -> None:
        conn = self._connection()
        async with self._write_lock:
            await conn.execute(sql, params)
            await conn.commit()