# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
"""Database introspection and query execution tool.

Provides SQL database access via SQLAlchemy. Supports any SQLAlchemy-compatible
database (PostgreSQL, MySQL, SQLite, Oracle, MS SQL Server, etc.).

Based on mcp-alchemy by Rui Machado (MPL 2.0).
https://github.com/runekaagaard/mcp-alchemy

Requires explicit db_url parameter for all operations.

Examples:
    # Get db_url from project config
    db.tables(db_url=proj.attr("myproject", "db_url"))

    # Or use literal URL
    db.tables(db_url="sqlite:///path/to/database.db")
    db.query("SELECT 1", db_url="postgresql://user:pass@localhost/dbname")
"""

from __future__ import annotations

# Pack for dot notation: db.tables(), db.schema(), db.query()
pack = "db"

__all__ = ["query", "schema", "tables"]

import contextlib
import threading
from collections import OrderedDict
from datetime import date, datetime
from typing import Any

from pydantic import BaseModel, Field
from sqlalchemy import Engine, create_engine, inspect, text

from ot.config import get_tool_config
from ot.logging import LogSpan
from ot.paths import resolve_cwd_path


class Config(BaseModel):
    """Pack configuration - discovered by registry."""

    max_chars: int = Field(
        default=4000,
        ge=100,
        le=100000,
        description="Maximum characters in query result output",
    )


def _get_config() -> Config:
    """Get db pack configuration."""
    return get_tool_config("db", Config)


# Connection pool keyed by URL - persists across calls in process
# Uses OrderedDict for LRU eviction with bounded size
_ENGINES_MAXSIZE = 8
_engines_lock = threading.Lock()
_engines: OrderedDict[str, Engine] = OrderedDict()


def _resolve_sqlite_url(db_url: str) -> str:
    """Resolve relative paths in SQLite URLs.

    SQLite URLs use the format sqlite:///path/to/db.
    If the path is relative, resolve it against the project working directory.

    Args:
        db_url: Database URL string

    Returns:
        URL with resolved path if SQLite with relative path, otherwise unchanged
    """
    if not db_url.startswith("sqlite:///"):
        return db_url

    # Extract path from sqlite:///path
    path = db_url[10:]  # len("sqlite:///") == 10

    # Skip in-memory databases and absolute paths
    if not path or path == ":memory:" or path.startswith("/"):
        return db_url

    # Resolve relative path against project directory
    resolved = resolve_cwd_path(path)
    return f"sqlite:///{resolved}"


def _create_engine(db_url: str) -> Engine:
    """Create SQLAlchemy engine with MCP-optimized settings."""
    # MCP-optimized defaults
    options: dict[str, Any] = {
        "isolation_level": "AUTOCOMMIT",
        "pool_pre_ping": True,  # Test connections before use
        "pool_size": 1,  # Single connection for MCP patterns
        "max_overflow": 2,  # Allow temporary burst capacity
        "pool_recycle": 3600,  # Refresh connections older than 1hr
    }

    return create_engine(db_url, **options)


def _get_engine(db_url: str) -> Engine:
    """Get or create engine for given URL with retry logic."""
    # Resolve relative paths in SQLite URLs
    resolved_url = _resolve_sqlite_url(db_url)

    # Fast path: check cache with lock
    with _engines_lock:
        if resolved_url in _engines:
            # LRU: move to end on access
            _engines.move_to_end(resolved_url)
            return _engines[resolved_url]

    # Create engine outside lock (slow operation)
    with LogSpan(span="db.connect", db_url=resolved_url) as span:
        try:
            engine = _create_engine(resolved_url)
        except Exception:
            span.add(retry=True)
            # One retry with fresh engine
            engine = _create_engine(resolved_url)

        # Double-check after acquiring lock
        with _engines_lock:
            if resolved_url in _engines:
                # Another thread created it while we were waiting
                engine.dispose()
                _engines.move_to_end(resolved_url)
                return _engines[resolved_url]

            _engines[resolved_url] = engine

            # LRU eviction: dispose oldest engine when over maxsize
            while len(_engines) > _ENGINES_MAXSIZE:
                _, oldest_engine = _engines.popitem(last=False)
                with contextlib.suppress(Exception):
                    oldest_engine.dispose()

            span.add(cached=False)
            return engine


def _format_value(val: Any) -> str:
    """Format a value for display."""
    if val is None:
        return "NULL"
    if isinstance(val, (datetime, date)):
        return val.isoformat()
    return str(val)


def tables(
    *, db_url: str, filter: str | None = None, ignore_case: bool = False
) -> str:
    """List table names in the database.

    Args:
        db_url: Database URL (required)
        filter: Optional substring to filter table names
        ignore_case: If True, filter matching is case-insensitive

    Returns:
        Comma-separated list of table names

    Example:
        # List all tables
        db.tables(db_url=proj.attr("myproject", "db_url"))

        # Filter tables containing "user"
        db.tables(db_url="sqlite:///data.db", filter="user")

        # Case-insensitive filter
        db.tables(db_url="sqlite:///data.db", filter="USER", ignore_case=True)
    """
    with LogSpan(span="db.tables", db_url=db_url, filter=filter) as s:
        if not db_url or not db_url.strip():
            s.add(error="empty_db_url")
            return "Error: db_url parameter is required"

        try:
            engine = _get_engine(db_url)
            with engine.connect() as conn:
                inspector = inspect(conn)
                all_tables = inspector.get_table_names()

                if filter:
                    if ignore_case:
                        filter_lower = filter.lower()
                        all_tables = [t for t in all_tables if filter_lower in t.lower()]
                    else:
                        all_tables = [t for t in all_tables if filter in t]

                s.add(resultCount=len(all_tables))
                return ", ".join(all_tables) if all_tables else "No tables found"

        except Exception as e:
            s.add(error=str(e))
            return f"Error: {e}"


def schema(*, table_names: list[str], db_url: str) -> str:
    """Get schema definitions for specified tables.

    Returns column names, types, primary keys, and foreign key relationships.

    Args:
        table_names: List of table names to inspect
        db_url: Database URL (required)

    Returns:
        Formatted schema definitions

    Example:
        # Single table
        db.schema(table_names=["users"], db_url=ot.project("myproject", attr="db_url"))

        # Multiple tables
        db.schema(table_names=["users", "orders"], db_url="sqlite:///data.db")
    """
    with LogSpan(span="db.schema", tables=table_names, db_url=db_url) as s:
        if not db_url or not db_url.strip():
            s.add(error="empty_db_url")
            return "Error: db_url parameter is required"

        if not table_names:
            s.add(error="no_tables")
            return "Error: table_names parameter is required"

        try:
            engine = _get_engine(db_url)
            with engine.connect() as conn:
                inspector = inspect(conn)
                results: list[str] = []

                for table_name in table_names:
                    results.append(_format_table_schema(inspector, table_name))

                s.add(resultCount=len(table_names))
                return "\n".join(results)

        except Exception as e:
            s.add(error=str(e))
            return f"Error: {e}"


def _format_table_schema(inspector: Any, table_name: str) -> str:
    """Format schema for a single table."""
    try:
        columns = inspector.get_columns(table_name)
    except Exception:
        return f"{table_name}: [table not found]"

    try:
        foreign_keys = inspector.get_foreign_keys(table_name)
        pk_constraint = inspector.get_pk_constraint(table_name)
    except Exception:
        foreign_keys = []
        pk_constraint = {}

    primary_keys = set(pk_constraint.get("constrained_columns", []))

    result = [f"{table_name}:"]

    # Process columns - use explicit key access to avoid mutating the dict
    show_key_only = {"nullable", "autoincrement"}
    skip_keys = {"name", "type", "comment"}
    for column in columns:
        name = column["name"]
        col_type = str(column["type"])

        parts = []
        if name in primary_keys:
            parts.append("primary key")
        parts.append(col_type)

        for k, v in column.items():
            if k in skip_keys:
                continue
            if v:
                if k in show_key_only:
                    parts.append(k)
                else:
                    parts.append(f"{k}={v}")

        result.append(f"    {name}: " + ", ".join(parts))

    # Process relationships
    if foreign_keys:
        result.extend(["", "    Relationships:"])
        for fk in foreign_keys:
            constrained = ", ".join(fk["constrained_columns"])
            referred_table = fk["referred_table"]
            referred_cols = ", ".join(fk["referred_columns"])
            result.append(f"      {constrained} -> {referred_table}.{referred_cols}")

    return "\n".join(result)


def query(*, sql: str, db_url: str, params: dict[str, Any] | None = None) -> str:
    """Execute a SQL query and return formatted results.

    IMPORTANT: Always use the params parameter for variable substitution
    (e.g., 'WHERE id = :id' with params={'id': 123}) to prevent SQL injection.

    Args:
        sql: SQL query to execute
        db_url: Database URL (required)
        params: Query parameters for safe substitution

    Returns:
        Formatted query results or error message

    Example:
        # Basic query
        db.query(sql="SELECT * FROM users LIMIT 5", db_url=ot.project("myproject", attr="db_url"))

        # Parameterized query (safe from SQL injection)
        db.query(
            sql="SELECT * FROM users WHERE status = :status",
            db_url="sqlite:///data.db",
            params={"status": "active"}
        )

        # INSERT/UPDATE/DELETE
        db.query(
            sql="UPDATE users SET status = :status WHERE id = :id",
            db_url="postgresql://user:pass@localhost/db",
            params={"status": "inactive", "id": 123}
        )
    """
    with LogSpan(span="db.query", sql=sql, db_url=db_url) as s:
        if not db_url or not db_url.strip():
            s.add(error="empty_db_url")
            return "Error: db_url parameter is required"

        if not sql or not sql.strip():
            s.add(error="empty_query")
            return "Error: sql parameter is required"

        try:
            engine = _get_engine(db_url)
            with engine.connect() as conn:
                cursor_result = conn.execute(text(sql), params or {})

                if not cursor_result.returns_rows:
                    affected = cursor_result.rowcount
                    s.add(rowsAffected=affected)
                    return f"Success: {affected} rows affected"

                max_chars = _get_config().max_chars
                output, row_count, truncated = _format_query_results(
                    cursor_result, max_chars
                )
                s.add(rows=row_count, truncated=truncated)
                return output

        except Exception as e:
            s.add(error=str(e))
            return f"Error: {e}"


def _format_query_results(cursor_result: Any, max_chars: int) -> tuple[str, int, bool]:
    """Format query results in vertical format.

    Args:
        cursor_result: SQLAlchemy cursor result
        max_chars: Maximum characters for output

    Returns:
        Tuple of (formatted_output, row_count, was_truncated)
    """
    result: list[str] = []
    size = 0
    row_count = 0
    displayed_count = 0
    truncated = False
    keys = list(cursor_result.keys())

    while row := cursor_result.fetchone():
        row_count += 1

        if truncated:
            continue

        sub_result = [f"{row_count}. row"]
        for col, val in zip(keys, row, strict=True):
            sub_result.append(f"{col}: {_format_value(val)}")
        sub_result.append("")

        row_size = sum(len(x) + 1 for x in sub_result)
        size += row_size

        if size > max_chars:
            truncated = True
        else:
            displayed_count += 1
            result.extend(sub_result)

    if row_count == 0:
        return "No rows returned", 0, False

    if truncated:
        result.append(
            f"Result: showing first {displayed_count} of {row_count} rows (output truncated)"
        )
    else:
        result.append(f"Result: {row_count} rows")

    return "\n".join(result), row_count, truncated
