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
    username: str | None
    timestamp: str
    prompt_text: str
    enhanced_prompt: str | None
    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

    @property
    def display_user(self) -> str:
        return f"@{self.username}" if self.username else str(self.user_id)

    @classmethod
    def from_row(cls, row: aiosqlite.Row) -> Request:
        return cls(
            id=row["id"],
            user_id=row["user_id"],
            username=row["username"],
            timestamp=row["timestamp"],
            prompt_text=row["prompt_text"],
            enhanced_prompt=row["enhanced_prompt"],
            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"],
        )


@dataclass(frozen=True)
class UserStat:
    user_id: int
    username: str | None
    total: int
    last_seen: str

    @property
    def display_user(self) -> str:
        return f"@{self.username}" if self.username else str(self.user_id)


SCHEMA = """
CREATE TABLE IF NOT EXISTS requests (
    id               INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id          INTEGER NOT NULL,
    username         TEXT,
    timestamp        TEXT    NOT NULL,
    prompt_text      TEXT    NOT NULL,
    enhanced_prompt  TEXT,
    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._migrate()
        await self._conn.commit()

    async def _migrate(self) -> None:
        conn = self._connection()
        async with conn.execute("PRAGMA table_info(requests)") as cursor:
            columns = {row["name"] async for row in cursor}
        if "username" not in columns:
            await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
        if "enhanced_prompt" not in columns:
            await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT")

    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,
        username: str | None = None,
        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, username, timestamp, prompt_text,
                         input_file_path, status, mode, seed)
                    VALUES (?, ?, ?, ?, ?, 'queued', ?, ?)
                    """,
                    (
                        user_id,
                        username,
                        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]

    @staticmethod
    def _filters(
        user_id: int | None,
        status: str | None,
        mode: str | None,
    ) -> tuple[str, list[Any]]:
        clauses: list[str] = []
        params: list[Any] = []
        if user_id is not None:
            clauses.append("user_id = ?")
            params.append(user_id)
        if status:
            clauses.append("status = ?")
            params.append(status)
        if mode:
            clauses.append("mode = ?")
            params.append(mode)
        where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
        return where, params

    async def list_requests(
        self,
        *,
        user_id: int | None = None,
        status: str | None = None,
        mode: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Request]:
        conn = self._connection()
        where, params = self._filters(user_id, status, mode)
        query = (
            f"SELECT * FROM requests{where} ORDER BY id DESC LIMIT ? OFFSET ?"
        )
        async with conn.execute(query, (*params, limit, offset)) as cursor:
            rows = await cursor.fetchall()
        return [Request.from_row(row) for row in rows]

    async def count_requests(
        self,
        *,
        user_id: int | None = None,
        status: str | None = None,
        mode: str | None = None,
    ) -> int:
        conn = self._connection()
        where, params = self._filters(user_id, status, mode)
        async with conn.execute(
            f"SELECT COUNT(*) AS total FROM requests{where}", params
        ) as cursor:
            row = await cursor.fetchone()
        return int(row["total"]) if row is not None else 0

    async def stats(self) -> dict[str, Any]:
        conn = self._connection()
        by_status: dict[str, int] = {}
        async with conn.execute(
            "SELECT status, COUNT(*) AS total FROM requests GROUP BY status"
        ) as cursor:
            async for row in cursor:
                by_status[str(row["status"])] = int(row["total"])
        async with conn.execute("SELECT COUNT(DISTINCT user_id) AS total FROM requests") as cursor:
            row = await cursor.fetchone()
        return {
            "total": sum(by_status.values()),
            "users": int(row["total"]) if row is not None else 0,
            "by_status": by_status,
        }

    async def list_users(self, limit: int = 100) -> list[UserStat]:
        conn = self._connection()
        async with conn.execute(
            """
            SELECT
                user_id,
                COUNT(*) AS total,
                MAX(timestamp) AS last_seen,
                (
                    SELECT username FROM requests AS r2
                    WHERE r2.user_id = requests.user_id AND r2.username IS NOT NULL
                    ORDER BY r2.id DESC LIMIT 1
                ) AS username
            FROM requests
            GROUP BY user_id
            ORDER BY last_seen DESC
            LIMIT ?
            """,
            (limit,),
        ) as cursor:
            rows = await cursor.fetchall()
        return [
            UserStat(
                user_id=int(row["user_id"]),
                username=row["username"],
                total=int(row["total"]),
                last_seen=str(row["last_seen"]),
            )
            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_enhanced_prompt(self, request_id: int, prompt: str) -> None:
        await self._update(
            "UPDATE requests SET enhanced_prompt = ? WHERE id = ?",
            (prompt, 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()