import asyncio
from contextlib import asynccontextmanager
from typing import Any, List, Optional

import aiosqlite
from loguru import logger

from core.config import Config


class DatabaseManager:
    def __init__(self, config: Config):
        self.db_path = config.database_config.path
        self._pool_size = config.database_config.pool_size
        self._pool: List[aiosqlite.Connection] = []
        self._pool_lock = asyncio.Lock()
        
    async def initialize(self):
        """初始化数据库连接池"""
        logger.info("初始化数据库连接池")
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute('''
                CREATE TABLE IF NOT EXISTS email_accounts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    email TEXT UNIQUE NOT NULL,
                    password TEXT NOT NULL,
                    client_id TEXT NOT NULL,
                    refresh_token TEXT NOT NULL,
                    in_use BOOLEAN DEFAULT 0,
                    cursor_password TEXT,
                    cursor_cookie TEXT,
                    sold BOOLEAN DEFAULT 0,
                    status TEXT DEFAULT 'pending',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            await db.commit()
            
        # 初始化连接池
        for i in range(self._pool_size):
            conn = await aiosqlite.connect(self.db_path)
            self._pool.append(conn)
        logger.info(f"数据库连接池初始化完成，大小: {self._pool_size}")
    
    async def cleanup(self):
        """清理数据库连接"""
        for conn in self._pool:
            await conn.close()
        self._pool.clear()
    
    @asynccontextmanager
    async def get_connection(self):
        """获取数据库连接"""
        async with self._pool_lock:
            if not self._pool:
                conn = await aiosqlite.connect(self.db_path)
            else:
                conn = self._pool.pop()
                
        try:
            yield conn
        finally:
            if len(self._pool) < self._pool_size:
                self._pool.append(conn)
            else:
                await conn.close()
    
    async def execute(self, query: str, params: tuple = ()) -> Any:
        """执行SQL语句"""
        async with self.get_connection() as conn:
            cursor = await conn.execute(query, params)
            await conn.commit()
            return cursor.lastrowid
                
    async def fetch_one(self, query: str, params: tuple = ()) -> Optional[tuple]:
        """查询单条记录"""
        async with self.get_connection() as conn:
            cursor = await conn.execute(query, params)
            return await cursor.fetchone()
                
    async def fetch_all(self, query: str, params: tuple = ()) -> List[tuple]:
        """查询多条记录"""
        async with self.get_connection() as conn:
            cursor = await conn.execute(query, params)
            return await cursor.fetchall()