import asyncio
import re
import time
from typing import Optional

from app.models.instance import Instance
from app.schemas.query import QueryResult
from app.services.sql_analyzer import SQLAnalyzer

MAX_ROWS = 10000


class QueryService:
    _analyzer = SQLAnalyzer()

    @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} 查询")

        # AST-based SQL validation (replaces keyword matching)
        is_safe, error_msg = QueryService._analyzer.validate_readonly(sql)
        if not is_safe:
            raise ValueError(error_msg)

        # Auto-add LIMIT if missing
        upper_sql = sql.strip().upper()
        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()