from __future__ import annotations

import argparse
import os
import sqlite3
import sys
from datetime import UTC, date, datetime
from pathlib import Path
from typing import Any

import psycopg
from psycopg.rows import dict_row

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))

from db import MIGRATIONS, SCHEMA_STATEMENTS, _normalize_dsn

REQUEST_COLUMNS = (
    "id",
    "user_id",
    "username",
    "timestamp",
    "prompt_text",
    "enhanced_prompt",
    "input_file_path",
    "output_file_path",
    "status",
    "comfy_prompt_id",
    "payment_charge_id",
    "payment_refunded",
    "is_premium",
    "mode",
    "seed",
    "error_text",
    "finished_at",
)

USER_COLUMNS = (
    "user_id",
    "username",
    "credits",
    "last_daily_credit",
    "referrer_id",
    "first_paid_at",
    "created_at",
    "is_admin",
    "is_blocked",
    "age_confirmed",
    "age_declined",
)

BLOCKED_COLUMNS = (
    "id",
    "user_id",
    "timestamp",
    "prompt",
    "reason",
    "mode",
    "has_person",
)

PAYMENT_COLUMNS = (
    "payload",
    "user_id",
    "timestamp",
    "mode",
    "prompt",
    "username",
    "input_file_path",
    "amount",
    "status",
    "is_premium",
    "telegram_payment_charge_id",
    "request_id",
)

SETTING_COLUMNS = ("key", "value")


def parse_dt(value: Any) -> datetime | None:
    if value in (None, ""):
        return None
    if isinstance(value, datetime):
        parsed = value
    else:
        parsed = datetime.fromisoformat(str(value))
    if parsed.tzinfo is None:
        parsed = parsed.replace(tzinfo=UTC)
    return parsed


def parse_day(value: Any) -> date | None:
    if value in (None, ""):
        return None
    if isinstance(value, datetime):
        return value.date()
    if isinstance(value, date):
        return value
    return date.fromisoformat(str(value))


def parse_bool(value: Any) -> bool:
    return bool(int(value or 0))


def table_names(connection: sqlite3.Connection) -> set[str]:
    rows = connection.execute(
        "SELECT name FROM sqlite_master WHERE type = 'table'"
    ).fetchall()
    return {str(row["name"]) for row in rows}


def read_rows(
    connection: sqlite3.Connection,
    table: str,
    columns: tuple[str, ...],
    defaults: dict[str, Any],
) -> list[dict[str, Any]]:
    if table not in table_names(connection):
        print(f"  {table}: в SQLite нет таблицы, пропуск")
        return []
    existing = {
        str(row["name"])
        for row in connection.execute(f"PRAGMA table_info({table})")
    }
    select = ", ".join(
        column if column in existing else f":{column} AS {column}"
        for column in columns
    )
    params = {
        column: defaults[column]
        for column in columns
        if column not in existing
    }
    query = f"SELECT {select} FROM {table}"
    cursor = (
        connection.execute(query, params)
        if params
        else connection.execute(query)
    )
    return [dict(row) for row in cursor.fetchall()]


def insert_rows(
    connection: psycopg.Connection,
    table: str,
    columns: tuple[str, ...],
    rows: list[dict[str, Any]],
    conflict: str | None,
) -> int:
    if not rows:
        return 0
    placeholders = ", ".join(["%s"] * len(columns))
    conflict_clause = (
        f"ON CONFLICT ({conflict}) DO NOTHING"
        if conflict
        else "ON CONFLICT DO NOTHING"
    )
    query = (
        f"INSERT INTO {table} ({', '.join(columns)}) "
        f"VALUES ({placeholders}) {conflict_clause}"
    )
    with connection.cursor() as cursor:
        cursor.executemany(
            query,
            [tuple(row[column] for column in columns) for row in rows],
        )
        return cursor.rowcount


def ensure_schema(connection: psycopg.Connection) -> None:
    with connection.cursor() as cursor:
        for statement in (*SCHEMA_STATEMENTS, *MIGRATIONS):
            cursor.execute(statement)


def reset_sequence(connection: psycopg.Connection, table: str) -> None:
    with connection.cursor() as cursor:
        cursor.execute(
            f"""
            SELECT setval(
                pg_get_serial_sequence('{table}', 'id'),
                COALESCE((SELECT MAX(id) FROM {table}), 1),
                (SELECT MAX(id) IS NOT NULL FROM {table})
            )
            """
        )


def migrate(sqlite_path: Path, database_url: str) -> None:
    if not sqlite_path.exists():
        raise SystemExit(f"SQLite-база не найдена: {sqlite_path}")

    source = sqlite3.connect(
        f"{sqlite_path.resolve().as_uri()}?mode=ro", uri=True
    )
    source.row_factory = sqlite3.Row
    try:
        users = read_rows(
            source,
            "users",
            USER_COLUMNS,
            {
                "username": None,
                "credits": 3,
                "last_daily_credit": date.today().isoformat(),
                "referrer_id": None,
                "first_paid_at": None,
                "created_at": datetime.now().astimezone().isoformat(),
                "is_admin": 0,
                "is_blocked": 0,
                "age_confirmed": 0,
                "age_declined": 0,
            },
        )
        requests = read_rows(
            source,
            "requests",
            REQUEST_COLUMNS,
            {
                "username": None,
                "enhanced_prompt": None,
                "payment_charge_id": None,
                "payment_refunded": 0,
                "is_premium": 0,
                "seed": None,
                "error_text": None,
                "finished_at": None,
            },
        )
        blocked = read_rows(
            source,
            "blocked_requests",
            BLOCKED_COLUMNS,
            {"mode": None, "has_person": 0},
        )
        payments = read_rows(
            source,
            "pending_payments",
            PAYMENT_COLUMNS,
            {"username": None, "input_file_path": None, "is_premium": 0},
        )
        settings = read_rows(source, "app_settings", SETTING_COLUMNS, {})
    finally:
        source.close()

    for row in users:
        row["last_daily_credit"] = parse_day(row["last_daily_credit"])
        row["first_paid_at"] = parse_dt(row["first_paid_at"])
        row["created_at"] = parse_dt(row["created_at"])
        row["is_admin"] = parse_bool(row["is_admin"])
        row["is_blocked"] = parse_bool(row["is_blocked"])
        row["age_confirmed"] = parse_bool(row["age_confirmed"])
        row["age_declined"] = parse_bool(row["age_declined"])

    for row in requests:
        row["timestamp"] = parse_dt(row["timestamp"])
        row["finished_at"] = parse_dt(row["finished_at"])
        row["payment_refunded"] = parse_bool(row["payment_refunded"])
        row["is_premium"] = parse_bool(row["is_premium"])

    for row in blocked:
        row["timestamp"] = parse_dt(row["timestamp"])
        row["has_person"] = parse_bool(row["has_person"])

    for row in payments:
        row["timestamp"] = parse_dt(row["timestamp"])
        row["is_premium"] = parse_bool(row["is_premium"])

    dsn = _normalize_dsn(database_url)
    with psycopg.connect(
        dsn, row_factory=dict_row, autocommit=True
    ) as target:
        ensure_schema(target)
        with target.transaction():
            inserted_users = insert_rows(
                target, "users", USER_COLUMNS, users, "user_id"
            )
            inserted_requests = insert_rows(
                target, "requests", REQUEST_COLUMNS, requests, None
            )
            inserted_blocked = insert_rows(
                target,
                "blocked_requests",
                BLOCKED_COLUMNS,
                blocked,
                None,
            )
            inserted_payments = insert_rows(
                target,
                "pending_payments",
                PAYMENT_COLUMNS,
                payments,
                "payload",
            )
            inserted_settings = insert_rows(
                target,
                "app_settings",
                SETTING_COLUMNS,
                settings,
                "key",
            )
            reset_sequence(target, "requests")
            reset_sequence(target, "blocked_requests")

    print(f"SQLite: {sqlite_path}")
    print("PostgreSQL: подключение выполнено")
    print(f"users:            {inserted_users} из {len(users)}")
    print(f"requests:         {inserted_requests} из {len(requests)}")
    print(f"blocked_requests: {inserted_blocked} из {len(blocked)}")
    print(f"pending_payments: {inserted_payments} из {len(payments)}")
    print(f"app_settings:     {inserted_settings} из {len(settings)}")
    print("Готово.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Перенос данных из SQLite в PostgreSQL"
    )
    parser.add_argument(
        "--sqlite",
        default=os.getenv("SQLITE_PATH", str(PROJECT_ROOT / "data" / "bot.db")),
        help="путь к старой SQLite-базе",
    )
    parser.add_argument(
        "--database-url",
        default=os.getenv("DATABASE_URL", ""),
        help="строка подключения PostgreSQL",
    )
    args = parser.parse_args()

    if not args.database_url:
        print(
            "DATABASE_URL не задан. Передай --database-url или переменную окружения.",
            file=sys.stderr,
        )
        raise SystemExit(1)

    migrate(Path(args.sqlite), args.database_url)


if __name__ == "__main__":
    main()