"""
Read-only DuckDB query runner with memory guardrails.

This module powers the `g-gremlin query` CLI command. It intentionally limits
SQL to read-only operations (SELECT/CTE/CREATE TEMP VIEW) so we do not imply
that CSV/Parquet files are mutable tables. Outputs are written via DuckDB
`COPY (SELECT ...) TO ...` for streaming, and stdout printing is capped for
small result sets to avoid accidental OOM.
"""

from __future__ import annotations

import csv
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import click

from g_gremlin.contracts import AgenticWarning
from g_gremlin.safety import SafetyContext

try:
    import duckdb  # type: ignore
except Exception:  # pragma: no cover - optional until command is invoked
    duckdb = None  # type: ignore

try:
    import psutil  # type: ignore
except Exception:  # pragma: no cover - optional dependency
    psutil = None


SUPPORTED_FORMATS = {"csv", "parquet", "json"}
STDOUT_ROW_LIMIT = 1000
DEFAULT_MEMORY_LIMIT = "2GB"


@dataclass
class QueryOutcome:
    """Outcome metadata for a query execution."""

    status: str
    row_count: Optional[int]
    output_path: Optional[str]
    output_format: Optional[str]
    warnings: list[str]
    explain_plan: Optional[str] = None


def _safe_memory_limit(requested: Optional[str]) -> str:
    """Return a safe DuckDB memory limit string."""
    if requested:
        return requested

    if psutil:
        try:
            total_ram = psutil.virtual_memory().total
            fifty_percent = total_ram // 2
            two_gb = 2 * 1024 * 1024 * 1024
            safe_limit_bytes = min(fifty_percent, two_gb)
            # Round down to whole GB for DuckDB consumption.
            safe_gb = max(1, safe_limit_bytes // (1024**3))
            return f"{safe_gb}GB"
        except Exception:
            pass

    return DEFAULT_MEMORY_LIMIT


def _ensure_temp_dir() -> Path:
    """Ensure a temp directory exists for DuckDB spill files."""
    temp_dir = Path(".g_gremlin") / "tmp" / "duckdb"
    temp_dir.mkdir(parents=True, exist_ok=True)
    return temp_dir


def _configure_duckdb(conn: duckdb.DuckDBPyConnection, memory_limit: Optional[str]) -> None:
    """Apply memory limit and temp directory settings to DuckDB."""
    limit = _safe_memory_limit(memory_limit)
    conn.execute(f"SET memory_limit='{limit}'")
    temp_dir = _ensure_temp_dir()
    conn.execute(f"SET temp_directory='{temp_dir}'")


def _detect_format(output: Optional[Path], explicit_format: Optional[str]) -> Optional[str]:
    if explicit_format:
        fmt = explicit_format.lower()
        if fmt not in SUPPORTED_FORMATS:
            raise click.ClickException(f"Unsupported format '{explicit_format}'. Choose from: csv, parquet, json.")
        return fmt

    if output:
        suffix = output.suffix.lower()
        if suffix == ".csv":
            return "csv"
        if suffix == ".parquet":
            return "parquet"
        if suffix == ".json":
            return "json"

    return None


def _enforce_read_only(sql: str) -> None:
    """Prevent mutating statements; allow SELECT/CTE/CREATE TEMP VIEW."""
    lowered = sql.strip().lower()

    allowed_prefixes = (
        lowered.startswith("select"),
        lowered.startswith("with"),
        lowered.startswith("describe"),
        lowered.startswith("show"),
        bool(re.match(r"^create\s+(or\s+replace\s+)?(temp|temporary)\s+view", lowered)),
    )
    if not any(allowed_prefixes):
        raise click.ClickException(
            "Only read-only queries are supported (SELECT/CTE/CREATE TEMP VIEW). "
            "Mutating statements like UPDATE/INSERT/DELETE are not allowed."
        )

    forbidden_patterns = [
        r"\binsert\b",
        r"\bupdate\b",
        r"\bdelete\b",
        r"\bmerge\b",
        r"\bdrop\b",
        r"\balter\b",
        r"\btruncate\b",
        r"\battach\b",
        r"\bcopy\s+into\b",
        r"\bcopy\s+from\b",
        r"\bcopy\s+to\b",
        r"\bcreate\s+table\b",
        r"\breplace\s+table\b",
    ]

    for pattern in forbidden_patterns:
        if re.search(pattern, lowered):
            raise click.ClickException(
                "Mutating statements are not supported in g-gremlin query. "
                "Use SELECT/CTE/CREATE TEMP VIEW only."
            )


def _validate_output_path(output: Path, overwrite: bool) -> None:
    if output.exists():
        if not overwrite:
            raise click.ClickException(
                f"Output file {output} exists. Pass --overwrite to replace it."
            )
        if output.is_file():
            output.unlink()
        else:
            raise click.ClickException(f"Output path {output} exists and is not a file.")
    if output.parent and not output.parent.exists():
        output.parent.mkdir(parents=True, exist_ok=True)


def _copy_to_output(conn: duckdb.DuckDBPyConnection, sql: str, output: Path, fmt: str) -> int:
    """Run COPY to materialize results; returns row count via COUNT(*) query."""
    if output.parent and not output.parent.exists():
        output.parent.mkdir(parents=True, exist_ok=True)
    fmt_upper = fmt.upper()
    output_escaped = str(output).replace("'", "''")
    copy_stmt = f"COPY ({sql}) TO '{output_escaped}' (FORMAT {fmt_upper}"
    if fmt == "csv":
        copy_stmt += ", HEADER"
    copy_stmt += ");"
    conn.execute(copy_stmt)

    # COUNT(*) can be expensive but gives a reliable summary. Fail softly on errors.
    row_count = None
    try:
        row_count = conn.execute(f"SELECT COUNT(*) FROM ({sql}) AS q").fetchone()[0]
    except Exception:
        pass
    return row_count if isinstance(row_count, int) else -1


def _stream_stdout(conn: duckdb.DuckDBPyConnection, sql: str, row_cap: int = STDOUT_ROW_LIMIT) -> int:
    """
    Print results to stdout as CSV with header. To avoid OOM, refuse to print
    more than `row_cap` rows; callers should route large queries to --output.
    """
    try:
        estimated = conn.execute(f"SELECT COUNT(*) FROM ({sql}) AS q").fetchone()[0]
    except Exception:
        estimated = None

    if estimated is not None and estimated > row_cap:
        raise click.ClickException(
            f"Query would return {estimated} rows; use --output to write to a file instead of stdout."
        )

    if estimated is None:
        rel = conn.execute(f"SELECT * FROM ({sql}) AS q LIMIT {row_cap + 1}")
    else:
        rel = conn.execute(sql)

    columns = [col[0] for col in rel.description or []]
    writer = csv.writer(sys.stdout)
    if columns:
        writer.writerow(columns)
    rows = rel.fetchall()
    if len(rows) > row_cap:
        raise click.ClickException(
            f"Query returned more than {row_cap} rows; use --output to stream results to a file."
        )
    writer.writerows(rows)
    return len(rows)


def run_query(
    sql: str,
    output: Optional[Path],
    overwrite: bool,
    explain: bool,
    fmt: Optional[str],
    memory_limit: Optional[str],
) -> QueryOutcome:
    """Execute the query and return metadata about the run."""
    if duckdb is None:
        raise click.ClickException(
            "duckdb is required for g-gremlin query. Install with `pip install duckdb`."
        )

    clean_sql = sql.strip().rstrip(";")
    _enforce_read_only(clean_sql)

    warnings: list[str] = []

    conn = duckdb.connect()
    _configure_duckdb(conn, memory_limit)
    outcome: QueryOutcome
    try:
        if explain:
            try:
                plan_rows = conn.execute(f"EXPLAIN {clean_sql}").fetchall()
                explain_plan = "\n".join(" ".join(str(part) for part in row) for row in plan_rows)
            except Exception as exc:
                raise click.ClickException(f"Failed to generate EXPLAIN plan: {exc}") from exc

            if output:
                warnings.append("Execution skipped because --explain was provided; no output file was written.")
            outcome = QueryOutcome(
                status="success",
                row_count=None,
                output_path=None,
                output_format=None,
                warnings=warnings,
                explain_plan=explain_plan,
            )
        else:
            row_count: Optional[int] = None
            output_path_str: Optional[str] = None
            output_format: Optional[str] = None

            if output:
                fmt_detected = _detect_format(output, fmt)
                if not fmt_detected:
                    raise click.ClickException(
                        "Could not determine output format. Provide --format or use a .csv/.parquet/.json extension."
                    )
                _validate_output_path(output, overwrite)
                output_format = fmt_detected
                row_count = _copy_to_output(conn, clean_sql, output, fmt_detected)
                output_path_str = str(output)
            else:
                row_count = _stream_stdout(conn, clean_sql)

            outcome = QueryOutcome(
                status="success",
                row_count=row_count if isinstance(row_count, int) and row_count >= 0 else None,
                output_path=output_path_str,
                output_format=output_format,
                warnings=warnings,
                explain_plan=None,
            )
    finally:
        conn.close()

    return outcome


def warning_to_agentic(message: str) -> AgenticWarning:
    """Convert a plain warning string to AgenticWarning."""
    return AgenticWarning(code="WARN", message=message, severity="warning")


def get_safety_context() -> SafetyContext:
    """Factory for SafetyContext used by the query command."""
    return SafetyContext(command="g-gremlin query", remote_mutation=False)
