from __future__ import annotations

from pathlib import Path

import aiosqlite

SCHEMA = """
CREATE TABLE IF NOT EXISTS users (
    user_id INTEGER PRIMARY KEY,
    selected_budget_id INTEGER,
    created_at TEXT NOT NULL DEFAULT (datetime('now'))
);

CREATE TABLE IF NOT EXISTS budgets (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    name TEXT NOT NULL DEFAULT '',
    total_amount REAL NOT NULL,
    start_date TEXT NOT NULL,
    end_date TEXT NOT NULL,
    currency TEXT NOT NULL DEFAULT 'RUB',
    is_active INTEGER NOT NULL DEFAULT 1,
    created_at TEXT NOT NULL DEFAULT (datetime('now')),
    FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS expenses (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id INTEGER NOT NULL,
    budget_id INTEGER NOT NULL,
    amount REAL NOT NULL,
    note TEXT,
    spent_at TEXT NOT NULL,
    created_at TEXT NOT NULL DEFAULT (datetime('now')),
    FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE,
    FOREIGN KEY (budget_id) REFERENCES budgets(id) ON DELETE CASCADE
);
"""


async def _ensure_indexes(conn: aiosqlite.Connection) -> None:
    await conn.executescript(
        """
        CREATE INDEX IF NOT EXISTS idx_budgets_user
            ON budgets(user_id, is_active, id);
        CREATE INDEX IF NOT EXISTS idx_expenses_budget_spent_at
            ON expenses(budget_id, spent_at);
        CREATE INDEX IF NOT EXISTS idx_expenses_user_spent_at
            ON expenses(user_id, spent_at);
        """
    )


async def _table_columns(conn: aiosqlite.Connection, table: str) -> set[str]:
    cursor = await conn.execute(f"PRAGMA table_info({table})")
    rows = await cursor.fetchall()
    return {str(row[1]) for row in rows}


async def _has_unique_user_on_budgets(conn: aiosqlite.Connection) -> bool:
    cursor = await conn.execute("PRAGMA index_list(budgets)")
    indexes = await cursor.fetchall()
    for idx in indexes:
        # (seq, name, unique, origin, partial)
        if not idx[2]:
            continue
        name = idx[1]
        info = await conn.execute(f"PRAGMA index_info({name})")
        cols = [row[2] for row in await info.fetchall()]
        if cols == ["user_id"]:
            return True
    return False


async def migrate_schema(conn: aiosqlite.Connection) -> None:
    """Upgrade legacy one-budget-per-user schema in place."""
    tables = {
        row[0]
        for row in await (
            await conn.execute(
                "SELECT name FROM sqlite_master WHERE type='table'"
            )
        ).fetchall()
    }
    if "budgets" not in tables:
        return

    budget_cols = await _table_columns(conn, "budgets")
    if "name" not in budget_cols:
        await conn.execute(
            "ALTER TABLE budgets ADD COLUMN name TEXT NOT NULL DEFAULT ''"
        )
    if "is_active" not in budget_cols:
        await conn.execute(
            "ALTER TABLE budgets ADD COLUMN is_active INTEGER NOT NULL DEFAULT 1"
        )

    user_cols = await _table_columns(conn, "users")
    if "selected_budget_id" not in user_cols:
        await conn.execute(
            "ALTER TABLE users ADD COLUMN selected_budget_id INTEGER"
        )

    if await _has_unique_user_on_budgets(conn):
        await conn.executescript(
            """
            CREATE TABLE budgets_migrated (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                name TEXT NOT NULL DEFAULT '',
                total_amount REAL NOT NULL,
                start_date TEXT NOT NULL,
                end_date TEXT NOT NULL,
                currency TEXT NOT NULL DEFAULT 'RUB',
                is_active INTEGER NOT NULL DEFAULT 1,
                created_at TEXT NOT NULL DEFAULT (datetime('now')),
                FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
            );
            INSERT INTO budgets_migrated (
                id, user_id, name, total_amount, start_date, end_date,
                currency, is_active, created_at
            )
            SELECT
                id, user_id,
                COALESCE(NULLIF(name, ''), 'Бюджет'),
                total_amount, start_date, end_date, currency,
                COALESCE(is_active, 1), created_at
            FROM budgets;
            DROP TABLE budgets;
            ALTER TABLE budgets_migrated RENAME TO budgets;
            """
        )

    expense_cols = await _table_columns(conn, "expenses")
    if "budget_id" not in expense_cols:
        await conn.execute("ALTER TABLE expenses ADD COLUMN budget_id INTEGER")
        await conn.execute(
            """
            UPDATE expenses
            SET budget_id = (
                SELECT b.id FROM budgets b
                WHERE b.user_id = expenses.user_id
                ORDER BY b.id DESC
                LIMIT 1
            )
            WHERE budget_id IS NULL
            """
        )
        # Drop orphan expenses that have no budget (should be rare)
        await conn.execute("DELETE FROM expenses WHERE budget_id IS NULL")
        await conn.executescript(
            """
            CREATE TABLE expenses_migrated (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                budget_id INTEGER NOT NULL,
                amount REAL NOT NULL,
                note TEXT,
                spent_at TEXT NOT NULL,
                created_at TEXT NOT NULL DEFAULT (datetime('now')),
                FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE,
                FOREIGN KEY (budget_id) REFERENCES budgets(id) ON DELETE CASCADE
            );
            INSERT INTO expenses_migrated (
                id, user_id, budget_id, amount, note, spent_at, created_at
            )
            SELECT id, user_id, budget_id, amount, note, spent_at, created_at
            FROM expenses;
            DROP TABLE expenses;
            ALTER TABLE expenses_migrated RENAME TO expenses;
            """
        )

    # Backfill selected budget for users who have budgets
    await conn.execute(
        """
        UPDATE users
        SET selected_budget_id = (
            SELECT b.id FROM budgets b
            WHERE b.user_id = users.user_id
            ORDER BY b.is_active DESC, b.id DESC
            LIMIT 1
        )
        WHERE selected_budget_id IS NULL
          AND EXISTS (SELECT 1 FROM budgets b WHERE b.user_id = users.user_id)
        """
    )
    await conn.execute(
        """
        UPDATE budgets
        SET name = 'Бюджет'
        WHERE name IS NULL OR TRIM(name) = ''
        """
    )
    await _ensure_indexes(conn)


class Database:
    def __init__(self, path: Path, *, read_only: bool = False) -> None:
        self.path = path
        self.read_only = read_only
        self._conn: aiosqlite.Connection | None = None

    async def connect(self) -> None:
        if self.read_only:
            if not self.path.exists():
                raise FileNotFoundError(
                    f"Database not found: {self.path}. "
                    "Сначала запусти API, чтобы создался data/budget.db"
                )
            uri = f"file:{self.path.resolve().as_posix()}?mode=ro"
            self._conn = await aiosqlite.connect(uri, uri=True)
            self._conn.row_factory = aiosqlite.Row
            await self._conn.execute("PRAGMA foreign_keys = ON")
            return

        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 foreign_keys = ON")
        try:
            await self._conn.execute("PRAGMA journal_mode=WAL")
        except aiosqlite.OperationalError:
            pass
        await self._conn.executescript(SCHEMA)
        await migrate_schema(self._conn)
        # Fresh DBs: migrate may no-op early if tables were just created with
        # full columns — still ensure indexes exist.
        budget_cols = await _table_columns(self._conn, "budgets")
        if "is_active" in budget_cols:
            await _ensure_indexes(self._conn)
        await self._conn.commit()

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

    @property
    def conn(self) -> aiosqlite.Connection:
        if self._conn is None:
            raise RuntimeError("Database is not connected")
        return self._conn