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
    payment_charge_id: str | None
    payment_refunded: bool
    is_premium: bool
    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"],
            payment_charge_id=row["payment_charge_id"],
            payment_refunded=bool(row["payment_refunded"]),
            is_premium=bool(row["is_premium"]),
            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)


@dataclass(frozen=True)
class BlockedRequest:
    id: int
    user_id: int
    timestamp: str
    prompt: str
    reason: str
    mode: str | None
    has_person: bool

    @classmethod
    def from_row(cls, row: aiosqlite.Row) -> BlockedRequest:
        return cls(
            id=row["id"],
            user_id=row["user_id"],
            timestamp=row["timestamp"],
            prompt=row["prompt"],
            reason=row["reason"],
            mode=row["mode"],
            has_person=bool(row["has_person"]),
        )


@dataclass(frozen=True)
class PendingPayment:
    payload: str
    user_id: int
    timestamp: str
    mode: str
    prompt: str
    username: str | None
    input_file_path: str | None
    amount: int
    status: str
    is_premium: bool
    telegram_payment_charge_id: str | None
    request_id: int | None

    @classmethod
    def from_row(cls, row: aiosqlite.Row) -> PendingPayment:
        return cls(
            payload=row["payload"],
            user_id=row["user_id"],
            timestamp=row["timestamp"],
            mode=row["mode"],
            prompt=row["prompt"],
            username=row["username"],
            input_file_path=row["input_file_path"],
            amount=int(row["amount"]),
            status=row["status"],
            is_premium=bool(row["is_premium"]),
            telegram_payment_charge_id=row["telegram_payment_charge_id"],
            request_id=row["request_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,
    payment_charge_id TEXT,
    payment_refunded INTEGER NOT NULL DEFAULT 0,
    is_premium       INTEGER NOT NULL DEFAULT 0,
    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');

CREATE TABLE IF NOT EXISTS blocked_requests (
    id         INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id    INTEGER NOT NULL,
    timestamp  TEXT    NOT NULL,
    prompt     TEXT    NOT NULL,
    reason     TEXT    NOT NULL,
    mode       TEXT,
    has_person INTEGER NOT NULL DEFAULT 0
);

CREATE INDEX IF NOT EXISTS idx_blocked_user_time
    ON blocked_requests(user_id, id);

CREATE TABLE IF NOT EXISTS pending_payments (
    payload                     TEXT PRIMARY KEY,
    user_id                     INTEGER NOT NULL,
    timestamp                   TEXT    NOT NULL,
    mode                        TEXT    NOT NULL CHECK (mode IN ('t2i', 'i2i')),
    prompt                      TEXT    NOT NULL,
    username                    TEXT,
    input_file_path             TEXT,
    amount                      INTEGER NOT NULL,
    status                      TEXT    NOT NULL DEFAULT 'pending'
                                CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
    is_premium                  INTEGER NOT NULL DEFAULT 0,
    telegram_payment_charge_id  TEXT,
    request_id                  INTEGER
);

CREATE INDEX IF NOT EXISTS idx_pending_user_time
    ON pending_payments(user_id, timestamp);
"""


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")
        if "payment_charge_id" not in columns:
            await conn.execute("ALTER TABLE requests ADD COLUMN payment_charge_id TEXT")
        if "payment_refunded" not in columns:
            await conn.execute(
                "ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0"
            )
        if "is_premium" not in columns:
            await conn.execute(
                "ALTER TABLE requests ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
            )
        async with conn.execute("PRAGMA table_info(pending_payments)") as cursor:
            pending_columns = {row["name"] async for row in cursor}
        if "is_premium" not in pending_columns:
            await conn.execute(
                "ALTER TABLE pending_payments ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
            )

    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,
        payment_charge_id: str | None = None,
        is_premium: bool = False,
    ) -> 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, payment_charge_id,
                         is_premium)
                    VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?, ?)
                    """,
                    (
                        user_id,
                        username,
                        utc_now(),
                        prompt_text,
                        input_file_path,
                        mode,
                        seed,
                        payment_charge_id,
                        int(is_premium),
                    ),
                )
                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()
        request = await self.get_request(request_id)
        if request is None:
            return 0
        if request.is_premium:
            query = """
                SELECT COUNT(*) AS total FROM requests
                WHERE is_premium = 1 AND status IN ('queued', 'running')
                  AND id < ?
            """
            params = (request_id,)
        else:
            query = """
                SELECT COUNT(*) AS total FROM requests
                WHERE status IN ('queued', 'running')
                  AND (is_premium = 1 OR (is_premium = 0 AND id < ?))
            """
            params = (request_id,)
        async with conn.execute(query, params) as cursor:
            row = await cursor.fetchone()
        return int(row["total"]) if row is not None else 0

    async def queue_counts(self) -> dict[str, int]:
        conn = self._connection()
        counts = {"premium": 0, "normal": 0}
        async with conn.execute(
            """
            SELECT is_premium, COUNT(*) AS total FROM requests
            WHERE status IN ('queued', 'running')
            GROUP BY is_premium
            """
        ) as cursor:
            async for row in cursor:
                key = "premium" if row["is_premium"] else "normal"
                counts[key] = int(row["total"])
        return counts

    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_refunded(self, request_id: int) -> None:
        await self._update(
            "UPDATE requests SET payment_refunded = 1 WHERE id = ?",
            (request_id,),
        )

    async def create_pending_payment(
        self,
        *,
        payload: str,
        user_id: int,
        mode: str,
        prompt: str,
        username: str | None,
        input_file_path: str | None,
        amount: int,
        is_premium: bool = False,
    ) -> PendingPayment:
        conn = self._connection()
        async with self._write_lock:
            await conn.execute(
                """
                INSERT INTO pending_payments
                    (payload, user_id, timestamp, mode, prompt, username,
                     input_file_path, amount, status, is_premium)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?)
                """,
                (
                    payload,
                    user_id,
                    utc_now(),
                    mode,
                    prompt,
                    username,
                    input_file_path,
                    amount,
                    int(is_premium),
                ),
            )
            await conn.commit()
        pending = await self.get_pending_payment(payload)
        if pending is None:
            raise RuntimeError("Не удалось прочитать счёт")
        return pending

    async def get_pending_payment(self, payload: str) -> PendingPayment | None:
        conn = self._connection()
        async with conn.execute(
            "SELECT * FROM pending_payments WHERE payload = ?", (payload,)
        ) as cursor:
            row = await cursor.fetchone()
        return PendingPayment.from_row(row) if row is not None else None

    async def set_pending_choice(
        self, payload: str, is_premium: bool, amount: int
    ) -> bool:
        conn = self._connection()
        async with self._write_lock:
            cursor = await conn.execute(
                """
                UPDATE pending_payments
                SET is_premium = ?, amount = ?
                WHERE payload = ? AND status = 'pending' AND amount = 0
                """,
                (int(is_premium), amount, payload),
            )
            await conn.commit()
            return cursor.rowcount > 0

    async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
        conn = self._connection()
        async with self._write_lock:
            cursor = await conn.execute(
                """
                UPDATE pending_payments
                SET status = 'paid', telegram_payment_charge_id = ?
                WHERE payload = ? AND status = 'pending'
                """,
                (charge_id, payload),
            )
            await conn.commit()
            return cursor.rowcount > 0

    async def mark_pending_refunded(self, payload: str) -> None:
        await self._update(
            "UPDATE pending_payments SET status = 'refunded' WHERE payload = ?",
            (payload,),
        )

    async def mark_pending_cancelled(self, payload: str) -> None:
        await self._update(
            """
            UPDATE pending_payments SET status = 'cancelled'
            WHERE payload = ? AND status = 'pending'
            """,
            (payload,),
        )

    async def attach_pending_request(self, payload: str, request_id: int) -> None:
        await self._update(
            "UPDATE pending_payments SET request_id = ? WHERE payload = ?",
            (request_id, payload),
        )

    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 log_blocked(
        self,
        *,
        user_id: int,
        prompt: str,
        reason: str,
        mode: str | None = None,
        has_person: bool = False,
    ) -> int:
        conn = self._connection()
        async with self._write_lock:
            cursor = await conn.execute(
                """
                INSERT INTO blocked_requests
                    (user_id, timestamp, prompt, reason, mode, has_person)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (user_id, utc_now(), prompt, reason, mode, int(has_person)),
            )
            await conn.commit()
            return int(cursor.lastrowid)

    async def list_blocked(self, limit: int = 100) -> list[BlockedRequest]:
        conn = self._connection()
        async with conn.execute(
            "SELECT * FROM blocked_requests ORDER BY id DESC LIMIT ?", (limit,)
        ) as cursor:
            rows = await cursor.fetchall()
        return [BlockedRequest.from_row(row) for row in rows]

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