import json
import re

from sqlalchemy import text

from pocket_analyst._session import _Session
from pocket_analyst._db import connection, schema

MAX_ROWS = 500
_DML = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE|REPLACE|MERGE)\b",
    re.IGNORECASE,
)


def list_tables(session: _Session) -> list[str]:
    engine = connection.get_engine(session)
    return schema.get_tables(engine)


def run_query(sql: str, session: _Session, max_rows: int = MAX_ROWS) -> dict:
    stripped = sql.strip().lstrip(";").strip()
    upper = stripped.upper()
    if not (upper.startswith("SELECT") or upper.startswith("WITH")):
        raise ValueError("Only SELECT queries are permitted.")
    if ";" in stripped:
        raise ValueError("Multiple statements are not permitted.")
    if _DML.search(stripped):
        raise ValueError("Only SELECT queries are permitted.")

    engine = connection.get_engine(session)
    with engine.connect() as conn:
        result = conn.execute(text(sql))
        columns = list(result.keys())
        rows = result.fetchmany(max_rows)

    data = [dict(zip(columns, row)) for row in rows]
    output = {"columns": columns, "rows": data, "row_count": len(data)}
    if len(rows) == max_rows:
        output["warning"] = f"Results truncated at {max_rows} rows."
    return output
