"""执行模型写的 SQL：连接只能读本次可见的 Parquet 文件，只允许单条 SELECT

每次查询新建一个内存连接：
- allowed_paths 只放这些 Parquet 文件；不用 allowed_directories，否则同目录下的其他文件能被 glob 读到，
  COPY、ATTACH 也能往目录里写
- 关闭外部访问、扩展加载和 Python 变量替换扫描，最后锁定配置，SQL 里的 SET 一律失败
"""

import math
import threading
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from pathlib import Path
from typing import Any

import duckdb

MEMORY_LIMIT = "512MB"
THREADS = 2
MAX_ROWS = 200
MAX_CELL_CHARS = 200


class QueryError(ValueError):
    """查询失败的原因，原样返回给模型修正"""


@dataclass(frozen=True)
class QueryResult:
    columns: list[str]
    rows: list[list[Any]]
    # 结果超过 max_rows 行，只返回了前 max_rows 行
    truncated: bool

    def to_dict(self) -> dict[str, Any]:
        return {
            "columns": self.columns,
            "rows": self.rows,
            "row_count": len(self.rows),
            "truncated": self.truncated,
        }


def run_query(
    sql: str,
    tables: Mapping[str, Path],
    *,
    timeout: float,
    max_rows: int = MAX_ROWS,
) -> QueryResult:
    """tables 是表名到 Parquet 文件的映射，SQL 里只能用这些表名"""
    try:
        connection = connect(tables)
    except duckdb.Error as error:
        # 任务执行期间文件被删除或过期清理
        raise QueryError(
            f"数据表无法读取，文件可能已被删除：{_first_line(error)}"
        ) from None
    try:
        _check_statement(connection, sql, tables)
        timer = threading.Timer(timeout, connection.interrupt)
        timer.start()
        try:
            connection.execute(sql)
            columns = [column[0] for column in connection.description]
            rows = connection.fetchmany(max_rows + 1)
        except duckdb.InterruptException:
            raise QueryError(
                f"查询超过 {timeout:g} 秒被中止，请缩小范围或先聚合"
            ) from None
        except duckdb.PermissionException:
            raise QueryError(f"只能查询上传的数据表：{_table_list(tables)}") from None
        except duckdb.Error as error:
            raise QueryError(
                f"{_first_line(error)}；可用的表：{_table_list(tables)}"
            ) from None
        finally:
            timer.cancel()
    finally:
        connection.close()
    return QueryResult(
        columns,
        [[json_value(value) for value in row] for row in rows[:max_rows]],
        len(rows) > max_rows,
    )


def connect(
    tables: Mapping[str, Path], *, threads: int = THREADS
) -> duckdb.DuckDBPyConnection:
    """建好视图并锁定的连接；列概况也用它，和模型查询看到的数据一致

    threads 只给解析子进程算列概况时放宽，模型写的查询一律用 THREADS
    """
    connection = duckdb.connect(
        ":memory:",
        config={
            "autoinstall_known_extensions": False,
            "autoload_known_extensions": False,
            "python_enable_replacements": False,
            "memory_limit": MEMORY_LIMIT,
            "threads": threads,
        },
    )
    try:
        paths = [str(Path(path).resolve()) for path in tables.values()]
        for name, path in zip(tables, paths):
            connection.execute(
                f"CREATE VIEW {quote_identifier(name)} AS FROM read_parquet({_literal(path)})"
            )
        if paths:
            connection.execute("SET allowed_paths = ?", [paths])
        connection.execute("SET enable_external_access = false")
        connection.execute("SET lock_configuration = true")
    except BaseException:
        connection.close()
        raise
    return connection


def quote_identifier(name: str) -> str:
    return '"' + name.replace('"', '""') + '"'


def json_value(value: Any, max_chars: int = MAX_CELL_CHARS) -> Any:
    """查询结果转成 JSON 能表示的值；长文本截断"""
    if value is None or isinstance(value, (bool, int)):
        return value
    if isinstance(value, float):
        return None if math.isnan(value) or math.isinf(value) else value
    if isinstance(value, Decimal):
        return float(value)
    if isinstance(value, str):
        return value if len(value) <= max_chars else value[:max_chars] + "…"
    if isinstance(value, (datetime, date, time)):
        return value.isoformat()
    if isinstance(value, (list, tuple)):
        return [json_value(item, max_chars) for item in value]
    if isinstance(value, dict):
        return {str(key): json_value(item, max_chars) for key, item in value.items()}
    if isinstance(value, timedelta):
        return str(value)
    return json_value(str(value), max_chars)


def _check_statement(
    connection: duckdb.DuckDBPyConnection, sql: str, tables: Mapping[str, Path]
) -> None:
    try:
        statements = connection.extract_statements(sql)
    except duckdb.Error as error:
        raise QueryError(
            f"{_first_line(error)}；可用的表：{_table_list(tables)}"
        ) from None
    # PIVOT 不写 IN 时会先建一个枚举类型，解析出来是两条语句
    if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
        raise QueryError(
            "只能执行一条 SELECT 查询（可以用 WITH、UNPIVOT；PIVOT 要用 IN 写明取值）"
        )


def _literal(text: str) -> str:
    return "'" + text.replace("'", "''") + "'"


def _first_line(error: Exception) -> str:
    return str(error).strip().splitlines()[0]


def _table_list(tables: Mapping[str, Path]) -> str:
    return "、".join(tables) if tables else "（无）"