"""Execute SQL queries against the dbt project's database.

This module implements the query_database tool for dbt Core MCP.
"""

import csv
import io
import json
import logging
import re
import tempfile
from pathlib import Path
from typing import Any

from fastmcp.dependencies import Depends  # type: ignore[reportAttributeAccessIssue]
from fastmcp.server.context import Context

from ..context import DbtCoreServerContext
from ..cte_generator import generate_cte_model
from ..dependencies import get_state
from . import dbtTool

logger = logging.getLogger(__name__)


def extract_cte_sql(
    project_dir: Path,
    cte_name: str,
    model_name: str,
    additional_sql: str = "",
    model_paths: list[str] | None = None,
) -> str:
    """Extract CTE SQL from a model file and optionally replace the final SELECT.

    This function extracts a specific CTE from a dbt model file, resolves all its
    upstream dependencies, and optionally replaces the final SELECT with a full query.

    Args:
        project_dir: Path to the dbt project directory
        cte_name: Name of the CTE to extract
        model_name: Name of the model file containing the CTE (without .sql extension)
        additional_sql: Optional full SELECT/WITH query to replace the final SELECT
            (use __cte__ or {{ cte }}; replaced with the CTE name at execution)
        model_paths: List of model directory paths (defaults to ["models"])

    Returns:
        Complete SQL ready to execute (either the extracted CTE or replaced with full SQL)

    Raises:
        ValueError: If model file not found, multiple files found, or CTE extraction fails

    Examples:
        # Extract a CTE with a full SELECT
        sql = extract_cte_sql(
            project_dir,
            "customer_agg",
            "customers",
            "SELECT * FROM __cte__ WHERE order_count > 5 LIMIT 10"
        )

        # Extract a CTE and replace final SELECT
        sql = extract_cte_sql(
            project_dir,
            "customer_agg",
            "customers",
            "SELECT customer_id, COUNT(*) AS cnt FROM __cte__ GROUP BY customer_id"
        )
    """
    # Use default model_paths if not provided
    if model_paths is None:
        model_paths = ["models"]

    # Find the model file - search all configured model paths
    model_files = []
    for model_path in model_paths:
        models_dir = project_dir / model_path
        if models_dir.exists():
            model_files.extend(list(models_dir.rglob(f"{model_name}.sql")))

    if not model_files:
        paths_searched = ", ".join(model_paths)
        raise ValueError(f"Model file '{model_name}.sql' not found in any model paths: {paths_searched}")

    if len(model_files) > 1:
        raise ValueError(f"Multiple model files found for '{model_name}': {[str(f) for f in model_files]}")

    model_file = model_files[0]
    logger.info(f"Extracting CTE '{cte_name}' from model '{model_name}' at {model_file}")

    # Create a temporary file for the extracted CTE model
    # Use system temp directory to avoid dbt picking it up as a model
    with tempfile.NamedTemporaryFile(mode="w", suffix=".sql", delete=False) as tmp_file:
        tmp_path = Path(tmp_file.name)

        try:
            # Generate CTE model using the generator logic
            # The generate_cte_model expects empty test_given list when we're just extracting
            success = generate_cte_model(
                base_model_path=model_file,
                cte_name=cte_name,
                test_given=[],  # No CTE mocking when querying
                output_path=tmp_path,
            )

            if not success:
                example = 'Example: query_database(cte_name="%s", model_name="%s", sql="SELECT * FROM __cte__ WHERE <filters>" )' % (cte_name, model_name)
                raise ValueError("Failed to extract CTE '%s' from model '%s'. Ensure the CTE name matches the model definition and include __cte__ in your SQL. %s" % (cte_name, model_name, example))

            # Read the generated CTE SQL
            cte_sql = tmp_path.read_text()

            # Remove the sqlfluff disable comment if present
            cte_sql = re.sub(r"^-- sqlfluff:disable\s*\n", "", cte_sql, flags=re.MULTILINE)

            # If user provided additional SQL, require a full SELECT/WITH query
            # The CTE extraction already includes "select * from {cte_name}" at the end
            additional_sql_stripped = additional_sql.strip() if additional_sql else ""
            if additional_sql_stripped:
                if not re.match(r"^(select|with)\b", additional_sql_stripped, flags=re.IGNORECASE):
                    # Construct the correct query from the partial SQL provided
                    suggested_query = f"SELECT * FROM __cte__ {additional_sql_stripped}"
                    raise ValueError(f"additional_sql must be a full SELECT/WITH query when querying a CTE. Did you mean: {suggested_query}")

                # Ensure the query references the CTE via __cte__ or {{ cte }}
                if not re.search(r"__cte__|{{\s*cte\s*}}", additional_sql_stripped, flags=re.IGNORECASE):
                    raise ValueError("additional_sql must reference the CTE using __cte__ or {{ cte }}. Example: SELECT * FROM __cte__ WHERE <condition>")

                pattern = rf"select \* from {re.escape(cte_name)}$"
                query_sql = re.sub(r"\{\{\s*cte\s*\}\}", cte_name, additional_sql_stripped, flags=re.IGNORECASE)
                query_sql = query_sql.replace("__cte__", cte_name)
                final_sql = re.sub(pattern, query_sql, cte_sql, flags=re.IGNORECASE | re.MULTILINE)
            else:
                # Use the CTE SQL as-is (already has select * from cte_name)
                final_sql = cte_sql

            logger.debug(f"Final SQL to execute:\n{final_sql[:500]}...")
            return final_sql

        finally:
            # Clean up temporary file
            # Use a small delay to allow processes to release the file
            import time

            time.sleep(0.1)
            try:
                if tmp_path.exists():
                    tmp_path.unlink()
            except PermissionError:
                # File is still in use (Windows), try to delete it later
                logger.debug(f"Temporary CTE file {tmp_path} still in use, will be cleaned up by OS")
            except Exception as e:
                logger.warning(f"Failed to delete temporary CTE file {tmp_path}: {e}")


async def _implementation(
    ctx: Context | None,
    sql: str,
    output_file: str | None,
    output_format: str,
    cte_name: str | None,
    model_name: str | None,
    state: DbtCoreServerContext,
) -> dict[str, Any]:
    """Implementation function for query_database tool.

    Separated for testing purposes - tests call this directly with explicit state.
    The @tool() decorated query_database() function calls this with injected dependencies.
    """
    # Ensure dbt components are initialized
    await state.ensure_initialized(ctx, force_parse=False)

    # Normalize output_file path to be relative to workspace root if relative
    if output_file:
        output_path_obj = Path(output_file)
        if not output_path_obj.is_absolute():
            # Relative path - make it relative to workspace root
            if state.project_dir:
                output_file = str(state.project_dir / output_path_obj)
                logger.debug(f"Normalized relative output path to: {output_file}")
            else:
                logger.warning("Project directory not available, using output_file path as-is")

    async def progress_callback(current: int, total: int, message: str) -> None:
        if ctx:
            await ctx.report_progress(progress=current, total=total, message=message)

    # Handle CTE query if cte_name is provided
    if cte_name:
        if not model_name:
            raise ValueError("model_name is required when querying a CTE (cte_name is specified)")

        if not state.project_dir:
            raise ValueError("Project directory not initialized")

        # Get configured model paths
        model_paths = state.get_project_paths()["model-paths"]

        # Extract CTE SQL using the dedicated function
        sql = extract_cte_sql(
            project_dir=state.project_dir,
            cte_name=cte_name,
            model_name=model_name,
            additional_sql=sql,
            model_paths=model_paths,
        )

    # Execute query using dbt show with --no-populate-cache for optimal performance
    runner = await state.get_runner()
    result = await runner.invoke_query(sql, progress_callback=progress_callback if ctx else None)  # type: ignore

    if not result.success:
        error_msg = str(result.exception) if result.exception else "Unknown error"
        # Include dbt output in error message for context
        full_error = error_msg

        # Try to extract error from stdout or stderr
        if result.stdout and "Error" in result.stdout:
            full_error = result.stdout
        elif hasattr(result, "stderr") and result.stderr:
            full_error = result.stderr

        logger.error(f"Query execution failed. Error: {error_msg}, stdout: {result.stdout if hasattr(result, 'stdout') else 'N/A'}")
        raise RuntimeError(f"Query execution failed: {full_error}")

    # Parse JSON output from dbt show (extract the "show" payload)
    output = result.stdout if hasattr(result, "stdout") else ""

    try:
        # dbt show --output json returns: {"show": [...rows...]}
        # Find the JSON object (look for {"show": pattern)
        json_match = re.search(r'\{\s*"show"\s*:\s*\[', output)
        if not json_match:
            return {
                "status": "failed",
                "error": "No JSON output found in dbt show response",
            }

        # Use JSONDecoder to parse just the first complete JSON object
        # This handles extra data after the JSON (like log lines)
        decoder = json.JSONDecoder()
        data, _ = decoder.raw_decode(output, json_match.start())

        if "show" in data:
            rows = data["show"]
            row_count = len(rows)

            # Handle different output formats
            # Get elapsed time from result if available
            elapsed_time = result.elapsed_time if hasattr(result, "elapsed_time") and result.elapsed_time is not None else None

            if output_format in ("csv", "tsv"):
                # Convert to CSV/TSV format
                delimiter = "\t" if output_format == "tsv" else ","
                csv_buffer = io.StringIO()

                if rows:
                    writer = csv.DictWriter(csv_buffer, fieldnames=rows[0].keys(), delimiter=delimiter)
                    writer.writeheader()
                    writer.writerows(rows)
                    csv_string = csv_buffer.getvalue()
                else:
                    csv_string = ""

                if output_file:
                    # Save to file
                    output_path = Path(output_file)
                    output_path.parent.mkdir(parents=True, exist_ok=True)

                    with open(output_path, "w", encoding="utf-8", newline="") as f:
                        f.write(csv_string)

                    # Get file size
                    file_size_bytes = output_path.stat().st_size
                    file_size_kb = file_size_bytes / 1024

                    response: dict[str, Any] = {
                        "status": "success",
                        "row_count": row_count,
                        "format": output_format,
                        "saved_to": str(output_path),
                        "file_size_kb": round(file_size_kb, 2),
                    }
                    if elapsed_time is not None:
                        response["elapsed_time"] = round(elapsed_time, 2)
                    return response
                else:
                    # Return CSV/TSV inline
                    response: dict[str, Any] = {
                        "status": "success",
                        "row_count": row_count,
                        "format": output_format,
                        output_format: csv_string,
                    }
                    if elapsed_time is not None:
                        response["elapsed_time"] = round(elapsed_time, 2)
                    return response
            else:
                # JSON format (default)
                # Get elapsed time from result if available
                elapsed_time = result.elapsed_time if hasattr(result, "elapsed_time") and result.elapsed_time is not None else None

                if output_file:
                    # Ensure directory exists
                    output_path = Path(output_file)
                    output_path.parent.mkdir(parents=True, exist_ok=True)

                    # Write rows to file
                    with open(output_path, "w", encoding="utf-8") as f:
                        json.dump(rows, f, indent=2)

                    # Get file size
                    file_size_bytes = output_path.stat().st_size
                    file_size_kb = file_size_bytes / 1024

                    # Return metadata with preview
                    response: dict[str, Any] = {
                        "status": "success",
                        "row_count": row_count,
                        "saved_to": str(output_path),
                        "file_size_kb": round(file_size_kb, 2),
                        "columns": list(rows[0].keys()) if rows else [],
                        "preview": rows[:3],  # First 3 rows as preview
                    }
                    if elapsed_time is not None:
                        response["elapsed_time"] = round(elapsed_time, 2)
                    return response
                else:
                    # Return all rows inline
                    response: dict[str, Any] = {
                        "status": "success",
                        "row_count": row_count,
                        "rows": rows,
                    }
                    if elapsed_time is not None:
                        response["elapsed_time"] = round(elapsed_time, 2)
                    return response
        else:
            return {
                "status": "failed",
                "error": "Unexpected JSON format from dbt show",
                "data": data,
            }

    except json.JSONDecodeError as e:
        return {
            "status": "error",
            "message": f"Failed to parse query results: {e}",
            "raw_output": output[:500],
        }


@dbtTool()
async def query_database(
    ctx: Context,
    sql: str,
    output_file: str | None = None,
    output_format: str = "json",
    cte_name: str | None = None,
    model_name: str | None = None,
    state: DbtCoreServerContext = Depends(get_state),
) -> dict[str, Any]:
    """Execute a SQL query against the dbt project's database.

    This tool compiles and runs SQL with Jinja templating support, allowing you to use
    {{ ref('model') }} and {{ source('src', 'table') }} in your queries.

    **SQL Templating**:
    - Use {{ ref('model_name') }} to reference dbt models
    - Use {{ source('source_name', 'table_name') }} to reference source tables
    - dbt compiles these to actual table names before execution

    **CTE Querying (LLM quick reference)**
    - Always pass `cte_name` + `model_name` in parameters (not in SQL)
    - Always write a normal `SELECT ... FROM __cte__ ...`
        - `__cte__` or `{{ cte }}` is replaced with the CTE name
    - What happens under the hood
        - Extracts the target CTE plus upstream CTEs from the model
        - Runs your query against that extracted CTE
    - Templating
        - dbt resolves all `{{ ref() }}` / `{{ source() }}` automatically; no manual table names
    - Invalid syntax to avoid
        - `{{ ref('model', cte='name') }}` does **not** exist; always use `cte_name` + `model_name`

    **Output Management**:
    - For large result sets (>100 rows), use output_file to save results
    - If output_file is omitted, all data returns inline (may consume large context)
    - output_file is automatically created with parent directories
    - Relative paths are resolved relative to the workspace root

    **Output Formats**:
    - json (default): Returns data as JSON array of objects
    - csv: Returns comma-separated values with header row
    - tsv: Returns tab-separated values with header row
    - CSV/TSV formats use proper quoting (only when necessary) and are Excel-compatible

    Args:
           sql: SQL query with Jinja templating: {{ ref('model') }}, {{ source('src', 'table') }}
               For exploratory queries, include LIMIT. For aggregations/counts, omit it.
               When using cte_name/model_name, provide a full `SELECT`/`WITH` query that
               selects from `__cte__` (or `{{ cte }}`), which is replaced with the CTE name.
        output_file: Optional file path to save results. Recommended for large result sets (>100 rows).
                    Relative paths are resolved relative to the workspace root.
                    If provided, only metadata is returned (no preview for CSV/TSV).
                    If omitted, all data is returned inline (may consume large context).
        output_format: Output format - "json" (default), "csv", or "tsv"
        cte_name: Optional CTE name to query from a model (requires model_name)
        model_name: Optional model name containing the CTE (required when cte_name is specified)
        state: Shared state object injected by FastMCP

    Returns:
        JSON inline: {"status": "success", "row_count": N, "rows": [...], "elapsed_time": X.XX}
        JSON file: {"status": "success", "row_count": N, "saved_to": "path", "preview": [...], "elapsed_time": X.XX}
        CSV/TSV inline: {"status": "success", "row_count": N, "format": "csv", "csv": "...", "elapsed_time": X.XX}
        CSV/TSV file: {"status": "success", "row_count": N, "format": "csv", "saved_to": "path", "elapsed_time": X.XX}

        Note: elapsed_time is in seconds and represents the total query execution time including compilation

    Raises:
        RuntimeError: If query execution fails
        ValueError: If invalid CTE/model parameters provided

    Examples:
        # Simple query with ref()
        query_database(sql="SELECT * FROM {{ ref('customers') }} LIMIT 10")

        # Query with source()
        query_database(sql="SELECT * FROM {{ source('jaffle_shop', 'orders') }} LIMIT 5")

        # Aggregation (no LIMIT needed)
        query_database(sql="SELECT COUNT(*) as total FROM {{ ref('customers') }}")

        # Query a specific CTE from a model
        query_database(
            cte_name="customer_agg",
            model_name="customers",
            sql="SELECT * FROM __cte__ LIMIT 10"
        )

        # Query a CTE with filtering
        query_database(
            cte_name="customer_agg",
            model_name="customers",
            sql="SELECT * FROM __cte__ WHERE order_count > 5 LIMIT 20"
        )

        # Query a CTE with aggregation (full SELECT)
        query_database(
            cte_name="customer_agg",
            model_name="customers",
            sql="SELECT customer_id, COUNT(*) AS cnt FROM __cte__ GROUP BY customer_id"
        )

        # WRONG - Do NOT use ref() with cte parameter (does not exist):
        # query_database(sql="SELECT * FROM {{ ref('model', cte='cte_name') }}")
        #
        # CORRECT - Use cte_name and model_name parameters instead:
        # query_database(cte_name="cte_name", model_name="model", sql="SELECT * FROM __cte__ LIMIT 10")

        # Save large results to file
        query_database(
            sql="SELECT * FROM {{ ref('orders') }}",
            output_file="temp_auto/orders_export.json"
        )

        # Export as CSV
        query_database(
            sql="SELECT * FROM {{ ref('customers') }}",
            output_file="temp_auto/customers.csv",
            output_format="csv"
        )
    """
    return await _implementation(ctx, sql, output_file, output_format, cte_name, model_name, state)
