import asyncio
import re
import time
from typing import Optional

from app.models.instance import Instance
from app.schemas.query import QueryResult


FORBIDDEN_KEYWORDS = ["DROP", "TRUNCATE", "DELETE", "UPDATE", "INSERT", "CREATE", "ALTER", "GRANT", "REVOKE"]
MAX_ROWS = 10000


class QueryService:
    @staticmethod
    async def execute_query(
        instance: Instance, database: Optional[str], sql: str, user_id: int
    ) -> QueryResult:
        if instance.db_type != "mysql":
            raise ValueError(f"暂不支持 {instance.db_type} 查询")

        upper_sql = sql.strip().upper()
        for kw in FORBIDDEN_KEYWORDS:
            if upper_sql.startswith(kw) or f" {kw} " in upper_sql:
                raise ValueError(f"查询模式禁止执行 {kw} 操作")

        if not re.search(r"LIMIT\s+\d+", upper_sql, re.IGNORECASE):
            sql = sql.rstrip(";") + f" LIMIT {MAX_ROWS};"

        import aiomysql
        conn = await aiomysql.connect(
            host=instance.host, port=instance.port,
            user=instance.username, password=instance.password,
            db=database or instance.database,
        )

        start = time.time()
        try:
            async with conn.cursor(aiomysql.DictCursor) as cur:
                await asyncio.wait_for(cur.execute(sql), timeout=30)
                rows = await cur.fetchall()
                columns = [desc[0] for desc in cur.description] if cur.description else []
            elapsed = int((time.time() - start) * 1000)

            return QueryResult(
                columns=columns,
                rows=rows,
                total=len(rows),
                elapsed_ms=elapsed,
                truncated=len(rows) >= MAX_ROWS,
            )
        except asyncio.TimeoutError:
            raise ValueError("查询超时（30秒）")
        finally:
            conn.close()

    @staticmethod
    async def execute_dml(instance: Instance, database: str, sql: str) -> str:
        if instance.db_type != "mysql":
            raise ValueError(f"暂不支持 {instance.db_type} 执行")

        import aiomysql
        conn = await aiomysql.connect(
            host=instance.host, port=instance.port,
            user=instance.username, password=instance.password,
            db=database,
            autocommit=True,
        )

        try:
            async with conn.cursor() as cur:
                await asyncio.wait_for(cur.execute(sql), timeout=60)
                affected = cur.rowcount
            return f"执行成功，影响 {affected} 行"
        except asyncio.TimeoutError:
            raise ValueError("执行超时（60秒）")
        finally:
            conn.close()