from __future__ import annotations

import logging
import re
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, AsyncGenerator, Sequence, Type, TypeVar, cast

import aiomysql
import aiosqlite
from pydantic import BaseModel

from beni import bpath
from beni.btype import XPath

from .bfunc import get_sql_placement, to_any
from .btype import Null


class BModel(BaseModel):
    _table_name: str = ''

    @classmethod
    def table_name(cls):
        if type(cls._table_name) is not str:
            class_name = cls.__name__
            result = [class_name[0].lower()]
            for char in class_name[1:]:
                if char.isupper():
                    result.extend(['_', char.lower()])
                else:
                    result.append(char)
            cls._table_name = ''.join(result)
        return cls._table_name


_BModel = TypeVar('_BModel', bound=BModel)
_T = TypeVar('_T')


class ConditionMaker:
    _statement = ''

    def __init__(self):
        self._ary: list[str] = []
        self._args: list[Any] = []

    def add(self, sql: str, *args: Any):
        self._ary.append(sql)
        self._args.extend(args)

    @property
    def args(self) -> list[Any]:
        return self._args

    def __str__(self) -> str:
        if self._ary:
            return f'{self._statement} ' + ' AND '.join([f'({x})' for x in self._ary])
        else:
            return ''


class WhereMaker(ConditionMaker):
    _statement = 'WHERE'


class HavingMaker(ConditionMaker):
    _statement = 'HAVING'


class OnMaker(ConditionMaker):
    _statement = 'ON'


# ============================================================

AioCursor = aiomysql.Cursor | aiosqlite.Cursor
AioConnection = aiomysql.Connection | aiosqlite.Connection


class DbCursor:
    def __init__(self, cursor: AioCursor):
        self.cursor = cursor

    @property
    def connection(self) -> AioConnection:
        return to_any(self.cursor.connection)

    @property
    def lastrowid(self):
        return self.cursor.lastrowid

    @property
    def rowcount(self):
        return self.cursor.rowcount

    _sql_format_re = re.compile(r'\s*\n\s*')

    def sql_format(self, sql: str) -> str:
        return self._sql_format_re.sub(' ', sql).strip()

    async def commit(self):
        await self.connection.commit()

    async def rollback(self):
        await self.connection.rollback()

    async def execute(self, sql: str, *args: Any):
        sql = self.sql_format(sql)
        return await self.cursor.execute(sql, args)

    async def get_tuple_one(self, sql: str, *args: Any) -> tuple[Any, ...] | None:
        "获取单行数据，以 tuple 形式返回"
        await self.execute(sql, *args)
        result = cast(tuple[Any], await self.cursor.fetchone())
        return result or None

    async def get_dict_one(self, sql: str, *args: Any) -> dict[str, Any] | None:
        "获取单行数据，以 dict 形式返回"
        result = await self.get_tuple_one(sql, *args)
        if result:
            columns = self._get_columns()
            result = {v: result[i] for i, v in enumerate(columns)}
            return result
        return None

    async def get_one(self, model_class: Type[_BModel], sql: str, *args: Any) -> _BModel | None:
        "获取单行数据，以 modelClass 实例返回"
        data = await self.get_dict_one(sql, *args)
        if data:
            return model_class(**data)
        return None

    async def get_tuple_list(self, sql: str, *args: Any) -> list[tuple[Any, ...]]:
        "获取多行数据，以 tuple 形式返回"
        await self.execute(sql, *args)
        result = cast(list[tuple[Any]], await self.cursor.fetchall())
        return result or []

    async def get_dict_list(self, sql: str, *args: Any) -> list[dict[str, Any]]:
        "获取多行数据，以 dict 形式返回"
        result = await self.get_tuple_list(sql, *args)
        if result:
            columns = self._get_columns()
            return [{v: row[i] for i, v in enumerate(columns)} for row in result]
        return []

    async def get_list(self, model_class: Type[_BModel], sql: str, *args: Any) -> list[_BModel]:
        "获取多行数据，以 modelClass 实例返回"
        result = await self.get_dict_list(sql, *args)
        return [model_class(**x) for x in result]

    async def get_value(self, value_class: Type[_T], sql: str, *args: Any) -> _T | None:
        "获取单个值"
        await self.execute(sql, *args)
        result: Sequence[Any] = to_any(await self.cursor.fetchone())
        return result[0] if result else None

    async def get_value_list(self, value_class: Type[_T], sql: str, *args: Any) -> list[_T] | None:
        "获取多个值"
        await self.execute(sql, *args)
        result: Sequence[Sequence[Any]] = to_any(await self.cursor.fetchall())
        return [x[0] for x in result] if result else []

    async def add_one(self, model: BModel, *, table_name: str = '') -> Any:
        columns: list[str] = []
        values: list[Any] = []
        for k, v in model.model_dump(exclude_unset=True).items():
            columns.append(f'`{k}`')
            values.append(v)
        table_name = table_name or model.__class__.table_name()
        sql = f'INSERT INTO `{table_name}` ( {",".join(columns)} ) VALUES %s'
        await self.execute(sql, values)
        return self.lastrowid

    async def save_one(self, model: BModel, *, table_name: str = ''):
        columns: list[str] = []
        values: list[Any] = []
        for k, v in model.model_dump(exclude_unset=True).items():
            columns.append(f'`{k}`')
            values.append(v)
        table_name = table_name or model.__class__.table_name()
        update_sql = ','.join([f'{x} = VALUES( {x} )' for x in columns])
        sql = f"""INSERT INTO `{table_name}` ( {','.join(columns)} ) VALUES %s ON DUPLICATE KEY UPDATE {update_sql}"""
        return await self.execute(sql, values)

    async def add_list(self, model_list: Sequence[BModel], *, table_name: str = ''):
        assert model_list, 'modelList 必须至少有一个元素'
        columns: list[str] = []
        values: list[Sequence[Any]] = []
        for k in model_list[0].model_dump(exclude_unset=True).keys():
            columns.append(k)
        for model in model_list:
            values.append([getattr(model, x) for x in columns])
        columns = [f'`{x}`' for x in columns]
        table_name = table_name or model_list[0].__class__.table_name()
        sql = f"""INSERT INTO `{table_name}` ( {','.join(columns)} ) VALUES {get_sql_placement(columns, '%s')}"""
        sql = self.sql_format(sql)
        return await self.cursor.executemany(sql, values)

    async def save_list(self, model_list: Sequence[BModel], *, table_name: str = ''):
        assert model_list, 'modelList 必须至少有一个元素'
        columns: list[str] = []
        values: list[Sequence[Any]] = []
        for k in model_list[0].model_dump(exclude_unset=True).keys():
            columns.append(k)
        for model in model_list:
            values.append([getattr(model, x) for x in columns])
        columns = [f'`{x}`' for x in columns]
        table_name = table_name or model_list[0].__class__.table_name()
        update_sql = ','.join([f'{x} = VALUES( {x} )' for x in columns])
        sql = f"""INSERT INTO `{table_name}` ( {','.join(columns)} ) VALUES {get_sql_placement(columns, '%s')} ON DUPLICATE KEY UPDATE {update_sql}"""
        sql = self.sql_format(sql)
        return await self.cursor.executemany(sql, values)

    def _get_columns(self) -> tuple[str, ...]:
        return tuple(x[0] for x in to_any(self.cursor.description))


class Db:
    @asynccontextmanager
    async def get_cursor(self) -> AsyncGenerator[DbCursor, None]:
        yield Null

    def as_mysql_db(self):
        return cast(MysqlDb, self)

    def as_sqlite_db(self):
        return cast(SqliteDb, self)


# ==============================================================================


class MysqlCursor(DbCursor):
    pass


@dataclass
class MysqlDb(Db):
    host: str
    port: int
    user: str
    password: str
    db: str = ''
    _pool: aiomysql.Pool = Null

    @asynccontextmanager
    async def get_cursor(self) -> AsyncGenerator[MysqlCursor, None]:
        is_echo = logging.getLogger().level is logging.DEBUG
        if not self._pool:
            self._pool = await aiomysql.create_pool(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                db=self.db or None,
                echo=is_echo,
            )
        async with cast(aiomysql.Connection, self._pool.acquire()) as connection:
            async with cast(aiomysql.Cursor, connection.cursor()) as cursor:
                is_any_exception = False
                try:
                    yield MysqlCursor(cursor)
                except Exception:
                    is_any_exception = True
                    raise
                finally:
                    if is_any_exception:
                        await connection.rollback()
                    else:
                        await connection.commit()

    async def close(self):
        if self._pool:
            self._pool.close()
            await self._pool.wait_closed()


# ==============================================================================


class SqliteCursor(DbCursor):
    async def execute(self, sql: str, *args: Any):
        sql = super().sql_format(sql)
        values = []
        if args:
            for value in args:
                if type(value) in [set, list, tuple]:
                    sql = sql.replace('%s', self._make_placement(value), 1)
                    values.extend(list(value))
                else:
                    sql = sql.replace('%s', '?', 1)
                    values.append(value)
        return await self.cursor.execute(sql, values)

    def _make_placement(self, value: Sequence[Any] | set[Any]) -> str:
        return f'( {", ".join(["?"] * len(value))} )'

    async def save_one(self, model: BModel, *, table_name: str = ''):
        columns: list[str] = []
        values: list[Any] = []
        for k, v in model.model_dump(exclude_unset=True).items():
            columns.append(f'`{k}`')
            values.append(v)
        table_name = table_name or model.__class__.table_name()
        sql = f"""
            INSERT OR REPLACE INTO
                `{table_name}` ( {','.join(columns)} )
            VALUES
                %s
        """
        return await self.execute(sql, values)

    async def save_list(self, model_list: Sequence[BModel], *, table_name: str = ''):
        assert model_list, 'modelList 必须至少有一个元素'

        # 遍历所有modelList里面所有的modeul_dumps(exclude_unset=True).keys()，取并集
        columns: list[str] = []
        for model in model_list:
            columns.extend(model.model_dump(exclude_unset=True).keys())
        columns = list(set(columns))
        values: list[Sequence[Any]] = []
        for model in model_list:
            values.append([getattr(model, x) for x in columns])
        columns = [f'`{x}`' for x in columns]
        table_name = table_name or model_list[0].__class__.table_name()
        sql = f'INSERT OR REPLACE INTO `{table_name}` ( {",".join(columns)} ) VALUES {self._make_placement(columns)}'
        sql = self.sql_format(sql)
        return await self.cursor.executemany(sql, values)


_is_inited_sqlite = False


def _init_sqlite():
    global _is_inited_sqlite
    if not _is_inited_sqlite:
        import sqlite3

        sqlite3.register_converter(
            'bool',
            lambda x: (
                x
                not in (
                    b'',
                    b'0',
                    # None, # 如果是None根本就不会进来，这里判断也没有意义
                )
            ),
        )
        _is_inited_sqlite = True


class SqliteDb:
    def __init__(self, file: XPath):
        self.file = file
        _init_sqlite()

    @asynccontextmanager
    async def get_cursor(self) -> AsyncGenerator[SqliteCursor, None]:
        bpath.make(bpath.get(self.file).parent)
        async with aiosqlite.connect(self.file) as connection:
            async with connection.cursor() as cursor:
                is_any_exception = False
                try:
                    yield SqliteCursor(cursor)
                except Exception:
                    is_any_exception = True
                    raise
                finally:
                    if is_any_exception:
                        await connection.rollback()
                    else:
                        await connection.commit()
