from __future__ import annotations

"""MySQL/MariaDB database layer for the Empire bot.

Production uses a pooled MySQL/MariaDB connection. Explicit SQLite mode is
retained only for emergency rollback/tests and is not the production default.
"""

import hashlib
import sqlite3
import threading
from contextlib import nullcontext
from pathlib import Path
from typing import Any

from countries import CONTINENTS


def hash_code(code: str) -> str:
    return hashlib.sha256(code.strip().encode("utf-8")).hexdigest()


MYSQL_SCHEMA = (
    """CREATE TABLE IF NOT EXISTS countries (
        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
        continent VARCHAR(120) NOT NULL,
        name VARCHAR(191) NOT NULL,
        code_hash CHAR(64) NULL,
        enabled TINYINT NOT NULL DEFAULT 1,
        PRIMARY KEY (id), UNIQUE KEY uq_countries_name (name)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci""",
    """CREATE TABLE IF NOT EXISTS sessions (
        user_id BIGINT NOT NULL,
        country VARCHAR(191) NOT NULL,
        display_name VARCHAR(255) NOT NULL,
        username VARCHAR(191) NULL,
        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
        PRIMARY KEY (user_id), KEY ix_sessions_country (country)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci""",
    """CREATE TABLE IF NOT EXISTS admins (
        user_id BIGINT NOT NULL,
        username VARCHAR(191) NULL,
        added_by BIGINT NULL,
        protected TINYINT NOT NULL DEFAULT 0,
        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (user_id)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci""",
    """CREATE TABLE IF NOT EXISTS settings (
        `key` VARCHAR(191) NOT NULL,
        `value` TEXT NULL,
        PRIMARY KEY (`key`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci""",
    """CREATE TABLE IF NOT EXISTS images (
        `key` VARCHAR(191) NOT NULL,
        file_id VARCHAR(512) NOT NULL,
        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
        PRIMARY KEY (`key`)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci""",
    """CREATE TABLE IF NOT EXISTS campaigns (
        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
        kind VARCHAR(64) NOT NULL,
        country VARCHAR(191) NOT NULL,
        origin TEXT NOT NULL,
        destination TEXT NOT NULL,
        target_country TEXT NOT NULL,
        stats TEXT NOT NULL,
        eta TEXT NOT NULL,
        vanguard TEXT NOT NULL,
        commander_name VARCHAR(255) NOT NULL,
        sender_user_id BIGINT NOT NULL,
        news_chat_id VARCHAR(191) NULL,
        news_message_id BIGINT NULL,
        report_chat_id VARCHAR(191) NULL,
        report_message_id BIGINT NULL,
        private_chat_id VARCHAR(191) NULL,
        private_message_id BIGINT NULL,
        status VARCHAR(32) NOT NULL DEFAULT 'active',
        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (id),
        KEY ix_campaigns_sender (sender_user_id),
        KEY ix_campaigns_status (status),
        KEY ix_campaigns_country (country)
    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci""",
)


SQLITE_SCHEMA = """
CREATE TABLE IF NOT EXISTS countries (
    id INTEGER PRIMARY KEY AUTOINCREMENT, continent TEXT NOT NULL,
    name TEXT NOT NULL UNIQUE, code_hash TEXT, enabled INTEGER NOT NULL DEFAULT 1
);
CREATE TABLE IF NOT EXISTS sessions (
    user_id INTEGER PRIMARY KEY, country TEXT NOT NULL, display_name TEXT NOT NULL,
    username TEXT, updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS admins (
    user_id INTEGER PRIMARY KEY, username TEXT, added_by INTEGER,
    protected INTEGER NOT NULL DEFAULT 0, created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS settings (key TEXT PRIMARY KEY, value TEXT);
CREATE TABLE IF NOT EXISTS images (
    key TEXT PRIMARY KEY, file_id TEXT NOT NULL,
    updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS campaigns (
    id INTEGER PRIMARY KEY AUTOINCREMENT, kind TEXT NOT NULL, country TEXT NOT NULL,
    origin TEXT NOT NULL, destination TEXT NOT NULL, target_country TEXT NOT NULL,
    stats TEXT NOT NULL, eta TEXT NOT NULL, vanguard TEXT NOT NULL,
    commander_name TEXT NOT NULL, sender_user_id INTEGER NOT NULL,
    news_chat_id INTEGER, news_message_id INTEGER, report_chat_id INTEGER,
    report_message_id INTEGER, private_chat_id INTEGER, private_message_id INTEGER,
    status TEXT NOT NULL DEFAULT 'active', created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
);
"""


class Database:
    """Synchronous facade with pooled MySQL connections.

    The Telegram network path remains asynchronous. Each DB operation uses a
    short-lived pooled connection, so one user's commit does not hold a global
    SQLite file lock for other users.
    """

    def __init__(self, settings_or_path: Any):
        self._lock = threading.RLock()
        self.engine = "sqlite"
        self.conn = None
        self.pool = None
        if isinstance(settings_or_path, str):
            self._init_sqlite(settings_or_path)
        elif getattr(settings_or_path, "db_engine", "mysql").lower() == "sqlite":
            self._init_sqlite(settings_or_path.db_path)
        else:
            self._init_mysql(settings_or_path)
        self.init_schema()
        self.seed_countries()

    def _init_sqlite(self, path: str) -> None:
        self.engine = "sqlite"
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(path, check_same_thread=False, timeout=30)
        self.conn.row_factory = sqlite3.Row
        self.conn.execute("PRAGMA journal_mode=WAL")
        self.conn.execute("PRAGMA synchronous=NORMAL")
        self.conn.execute("PRAGMA busy_timeout=30000")

    def _init_mysql(self, settings: Any) -> None:
        try:
            from mysql.connector import pooling
        except ImportError as exc:
            raise RuntimeError(
                "mysql-connector-python is required when DB_ENGINE=mysql"
            ) from exc
        try:
            self.pool = pooling.MySQLConnectionPool(
                pool_name="empire_pool",
                pool_size=max(2, int(getattr(settings, "db_pool_size", 8))),
                pool_reset_session=True,
                host=settings.db_host,
                port=int(settings.db_port),
                database=settings.db_name,
                user=settings.db_user,
                password=settings.db_password,
                connection_timeout=int(getattr(settings, "db_connect_timeout", 10)),
                charset="utf8mb4",
                collation="utf8mb4_unicode_ci",
                autocommit=False,
            )
            self.engine = "mysql"
        except Exception as exc:
            raise RuntimeError(
                "Could not connect to MySQL/MariaDB. Check DB_HOST, DB_PORT, "
                "DB_NAME, DB_USER and DB_PASSWORD in .env."
            ) from exc

    @staticmethod
    def _mysql_sql(query: str) -> str:
        # Keeps the one legacy handler using SQLite '?' placeholders portable.
        return query.replace("?", "%s")

    def _guard(self):
        """Serialize the single SQLite connection, never the MySQL pool.

        MySQL connections are independent and thread-safe through the pool;
        putting a process-wide lock around them would turn the pool back into
        a single-file bottleneck under concurrent Telegram updates.
        """
        return self._lock if self.engine == "sqlite" else nullcontext()

    def _mysql_cursor(self, dictionary: bool = False):
        assert self.pool is not None
        conn = self.pool.get_connection()
        return conn, conn.cursor(dictionary=dictionary)

    def _mysql_read(self, query: str, params: tuple[Any, ...] = (), many: bool = False):
        conn, cursor = self._mysql_cursor(dictionary=True)
        try:
            cursor.execute(self._mysql_sql(query), params)
            return cursor.fetchall() if many else cursor.fetchone()
        finally:
            cursor.close(); conn.close()

    def _mysql_write(self, query: str, params: tuple[Any, ...] = (), *, last_id: bool = False):
        conn, cursor = self._mysql_cursor()
        try:
            cursor.execute(self._mysql_sql(query), params)
            conn.commit()
            return int(cursor.lastrowid or 0) if last_id else cursor.rowcount
        except Exception:
            conn.rollback(); raise
        finally:
            cursor.close(); conn.close()

    def init_schema(self) -> None:
        with self._guard():
            if self.engine == "mysql":
                for statement in MYSQL_SCHEMA:
                    self._mysql_write(statement)
                return
            assert self.conn is not None
            self.conn.executescript(SQLITE_SCHEMA)
            self._ensure_sqlite_column("sessions", "username", "TEXT")
            self._ensure_sqlite_column("campaigns", "news_chat_id", "INTEGER")
            self._ensure_sqlite_column("campaigns", "news_message_id", "INTEGER")
            self.conn.commit()

    def _ensure_sqlite_column(self, table: str, column: str, definition: str) -> None:
        assert self.conn is not None
        cols = {row[1] for row in self.conn.execute(f"PRAGMA table_info({table})")}
        if column not in cols:
            self.conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")

    def seed_countries(self) -> None:
        rows = [(continent, name) for continent, names in CONTINENTS.items() for name in names]
        with self._guard():
            if self.engine == "mysql":
                conn, cursor = self._mysql_cursor()
                try:
                    cursor.executemany("INSERT IGNORE INTO countries(continent,name) VALUES(%s,%s)", rows)
                    conn.commit()
                finally:
                    cursor.close(); conn.close()
                return
            assert self.conn is not None
            self.conn.executemany("INSERT OR IGNORE INTO countries(continent,name) VALUES(?,?)", rows)
            self.conn.commit()

    def rows(self, query: str, params: tuple[Any, ...] = ()) -> list[Any]:
        with self._guard():
            if self.engine == "mysql":
                return list(self._mysql_read(query, params, many=True))
            assert self.conn is not None
            return list(self.conn.execute(query, params).fetchall())

    def row(self, query: str, params: tuple[Any, ...] = ()) -> Any | None:
        with self._guard():
            if self.engine == "mysql":
                return self._mysql_read(query, params)
            assert self.conn is not None
            return self.conn.execute(query, params).fetchone()

    def _write_simple(self, query: str, params: tuple[Any, ...] = (), *, mysql_sql: str | None = None, last_id: bool = False) -> int:
        with self._guard():
            if self.engine == "mysql":
                return int(self._mysql_write(mysql_sql or query, params, last_id=last_id))
            assert self.conn is not None
            cur = self.conn.execute(query, params)
            self.conn.commit()
            return int(cur.lastrowid if last_id else cur.rowcount)

    def set_setting(self, key: str, value: str | None) -> None:
        self._write_simple(
            "INSERT INTO settings(`key`,`value`) VALUES(?,?) ON CONFLICT(`key`) DO UPDATE SET `value`=excluded.`value`",
            (key, value),
            mysql_sql="INSERT INTO settings(`key`,`value`) VALUES(?,?) ON DUPLICATE KEY UPDATE `value`=VALUES(`value`)",
        )

    def get_setting(self, key: str, default: str | None = None) -> str | None:
        row = self.row("SELECT `value` FROM settings WHERE `key`=?", (key,))
        return row["value"] if row and row["value"] is not None else default

    def set_code(self, country: str, code: str) -> None:
        code = code.strip()
        if not code:
            raise ValueError("Country code cannot be empty")
        self._write_simple("UPDATE countries SET code_hash=? WHERE name=?", (hash_code(code), country))

    def reset_codes(self) -> int:
        return self._write_simple("UPDATE countries SET code_hash=NULL")

    def authenticate(self, code: str) -> str | None:
        row = self.row("SELECT name FROM countries WHERE enabled=1 AND code_hash=?", (hash_code(code),))
        return row["name"] if row else None

    def set_session(self, user_id: int, country: str, display_name: str, username: str | None = None) -> None:
        self._write_simple(
            "INSERT INTO sessions(user_id,country,display_name,username) VALUES(?,?,?,?) ON CONFLICT(user_id) DO UPDATE SET country=excluded.country,display_name=excluded.display_name,username=excluded.username,updated_at=CURRENT_TIMESTAMP",
            (user_id, country, display_name, username),
            mysql_sql="INSERT INTO sessions(user_id,country,display_name,username) VALUES(?,?,?,?) ON DUPLICATE KEY UPDATE country=VALUES(country),display_name=VALUES(display_name),username=VALUES(username),updated_at=CURRENT_TIMESTAMP",
        )

    def get_session(self, user_id: int) -> Any | None:
        return self.row("SELECT * FROM sessions WHERE user_id=?", (user_id,))

    def clear_session(self, user_id: int) -> None:
        self._write_simple("DELETE FROM sessions WHERE user_id=?", (user_id,))

    def add_admin(self, user_id: int, username: str | None, added_by: int | None, protected: bool = False) -> None:
        username = username.lstrip("@").strip() if username else None
        self._write_simple(
            "INSERT INTO admins(user_id,username,added_by,protected) VALUES(?,?,?,?) ON CONFLICT(user_id) DO UPDATE SET username=COALESCE(excluded.username,admins.username),protected=MAX(admins.protected,excluded.protected)",
            (user_id, username, added_by, int(protected)),
            mysql_sql="INSERT INTO admins(user_id,username,added_by,protected) VALUES(?,?,?,?) ON DUPLICATE KEY UPDATE username=COALESCE(VALUES(username),admins.username),protected=GREATEST(admins.protected,VALUES(protected))",
        )

    def is_admin(self, user_id: int) -> bool:
        return self.row("SELECT 1 FROM admins WHERE user_id=?", (user_id,)) is not None

    def list_admins(self) -> list[Any]:
        return self.rows("SELECT * FROM admins ORDER BY protected DESC,created_at,user_id")

    def remove_admin(self, user_id: int) -> bool:
        return bool(self._write_simple("DELETE FROM admins WHERE user_id=? AND protected=0", (user_id,)))

    def find_user_by_username(self, username: str) -> Any | None:
        return self.row("SELECT user_id,username,display_name FROM sessions WHERE LOWER(username)=LOWER(?)", (username.lstrip("@"),))

    def set_image(self, key: str, file_id: str) -> None:
        self._write_simple(
            "INSERT INTO images(`key`,file_id) VALUES(?,?) ON CONFLICT(`key`) DO UPDATE SET file_id=excluded.file_id,updated_at=CURRENT_TIMESTAMP",
            (key, file_id),
            mysql_sql="INSERT INTO images(`key`,file_id) VALUES(?,?) ON DUPLICATE KEY UPDATE file_id=VALUES(file_id),updated_at=CURRENT_TIMESTAMP",
        )

    def get_image(self, key: str) -> str | None:
        row = self.row("SELECT file_id FROM images WHERE `key`=?", (key,))
        return row["file_id"] if row else None

    def add_image_rows(self, image_rows: list[tuple[str, str]], *, overwrite: bool = True) -> int:
        """Insert/update image mappings only; never touches country hashes."""
        if not image_rows:
            return 0
        with self._guard():
            if self.engine == "mysql":
                conn, cursor = self._mysql_cursor()
                try:
                    sql = "INSERT INTO images(`key`,file_id) VALUES(%s,%s)"
                    if overwrite:
                        sql += " ON DUPLICATE KEY UPDATE file_id=VALUES(file_id),updated_at=CURRENT_TIMESTAMP"
                    else:
                        sql = "INSERT IGNORE INTO images(`key`,file_id) VALUES(%s,%s)"
                    cursor.executemany(sql, image_rows)
                    conn.commit()
                    return len(image_rows)
                except Exception:
                    conn.rollback(); raise
                finally:
                    cursor.close(); conn.close()
            assert self.conn is not None
            sql = "INSERT INTO images(key,file_id) VALUES(?,?) ON CONFLICT(key) DO UPDATE SET file_id=excluded.file_id,updated_at=CURRENT_TIMESTAMP"
            if not overwrite:
                sql = "INSERT OR IGNORE INTO images(key,file_id) VALUES(?,?)"
            self.conn.executemany(sql, image_rows)
            self.conn.commit()
            return len(image_rows)

    def add_campaign(self, data: dict[str, Any]) -> int:
        cols = ("kind", "country", "origin", "destination", "target_country", "stats", "eta", "vanguard", "commander_name", "sender_user_id", "news_chat_id", "news_message_id", "report_chat_id", "report_message_id", "private_chat_id", "private_message_id")
        values = tuple(data.get(c) for c in cols)
        marks = ",".join("?" for _ in cols)
        return self._write_simple(f"INSERT INTO campaigns({','.join(cols)}) VALUES({marks})", values, last_id=True)

    def update_campaign(self, campaign_id: int, **fields: Any) -> None:
        allowed = {"news_chat_id", "news_message_id", "report_chat_id", "report_message_id", "private_chat_id", "private_message_id", "status"}
        fields = {k: v for k, v in fields.items() if k in allowed}
        if fields:
            sets = ",".join(f"{k}=?" for k in fields)
            self._write_simple(f"UPDATE campaigns SET {sets} WHERE id=?", (*fields.values(), campaign_id))

    def claim_campaign(self, campaign_id: int, expected_status: str, new_status: str) -> bool:
        return bool(self._write_simple("UPDATE campaigns SET status=? WHERE id=? AND status=?", (new_status, campaign_id, expected_status)))

    def get_campaign(self, campaign_id: int) -> Any | None:
        return self.row("SELECT * FROM campaigns WHERE id=?", (campaign_id,))

    def active_campaigns(self, limit: int = 20) -> list[Any]:
        return self.rows("SELECT * FROM campaigns WHERE status='active' ORDER BY id DESC LIMIT ?", (limit,))

    def close(self) -> None:
        if self.engine == "sqlite" and self.conn is not None:
            self.conn.close(); self.conn = None