"""
Database Tools

Direct database operations for agents:
- SQL query execution
- Table management
- Schema operations
- Data profiling
"""

from __future__ import annotations

import logging
import os
from datetime import datetime, timezone
from typing import Any

logger = logging.getLogger(__name__)


class DatabaseTools:
    """Tools for database operations"""

    def __init__(self, db_url: str | None = None):
        """
        Initialize database tools.

        Args:
            db_url: Database connection URL (PostgreSQL/TimescaleDB)
        """
        self.db_url = db_url or os.environ.get(
            "TIMESCALEDB_URL",
            "postgresql://localhost:5432/postgres"
        )
        self._pool = None
        logger.info("DatabaseTools initialized")

    async def initialize(self) -> None:
        """Initialize connection pool."""
        try:
            import asyncpg
            # Convert SQLAlchemy URL to asyncpg format
            url = self.db_url.replace("postgresql+asyncpg://", "postgresql://")
            self._pool = await asyncpg.create_pool(url, min_size=1, max_size=5)
            logger.info("Database connection pool created")
        except Exception as e:
            logger.warning(f"Failed to create connection pool: {e}")

    async def close(self) -> None:
        """Close connection pool."""
        if self._pool:
            await self._pool.close()

    async def execute_query(
        self,
        query: str,
        params: list | None = None,
        fetch: bool = True,
    ) -> dict[str, Any]:
        """
        Execute a SQL query.

        Args:
            query: SQL query to execute
            params: Query parameters
            fetch: Whether to fetch results

        Returns:
            Query results or execution status
        """
        if not self._pool:
            await self.initialize()

        if not self._pool:
            return {"error": "Database connection not available"}

        try:
            async with self._pool.acquire() as conn:
                if fetch:
                    rows = await conn.fetch(query, *(params or []))
                    return {
                        "success": True,
                        "rows": [dict(row) for row in rows],
                        "count": len(rows),
                    }
                else:
                    result = await conn.execute(query, *(params or []))
                    return {
                        "success": True,
                        "result": result,
                    }
        except Exception as e:
            return {"success": False, "error": str(e)}

    async def get_table_info(self, table_name: str, schema: str = "public") -> dict[str, Any]:
        """Get information about a table."""
        query = """
        SELECT
            column_name,
            data_type,
            is_nullable,
            column_default
        FROM information_schema.columns
        WHERE table_schema = $1 AND table_name = $2
        ORDER BY ordinal_position
        """
        result = await self.execute_query(query, [schema, table_name])
        if not result.get("success"):
            return result

        return {
            "table_name": table_name,
            "schema": schema,
            "columns": result.get("rows", []),
            "column_count": len(result.get("rows", [])),
        }

    async def get_table_stats(self, table_name: str, schema: str = "public") -> dict[str, Any]:
        """Get statistics for a table."""
        # Row count
        count_query = f'SELECT COUNT(*) as count FROM "{schema}"."{table_name}"'
        count_result = await self.execute_query(count_query)

        # Table size (PostgreSQL specific)
        size_query = f"SELECT pg_size_pretty(pg_total_relation_size('\"{schema}\".\"{table_name}\"')) as size"
        size_result = await self.execute_query(size_query)

        return {
            "table_name": table_name,
            "schema": schema,
            "row_count": count_result.get("rows", [{}])[0].get("count", 0) if count_result.get("success") else None,
            "size": size_result.get("rows", [{}])[0].get("size", "unknown") if size_result.get("success") else None,
        }

    async def list_tables(self, schema: str = "public") -> dict[str, Any]:
        """List all tables in a schema."""
        query = """
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = $1 AND table_type = 'BASE TABLE'
        ORDER BY table_name
        """
        result = await self.execute_query(query, [schema])
        if not result.get("success"):
            return result

        return {
            "schema": schema,
            "tables": [row["table_name"] for row in result.get("rows", [])],
            "count": len(result.get("rows", [])),
        }

    async def sample_data(
        self,
        table_name: str,
        schema: str = "public",
        limit: int = 10,
    ) -> dict[str, Any]:
        """Get a sample of data from a table."""
        query = f'SELECT * FROM "{schema}"."{table_name}" LIMIT {limit}'
        result = await self.execute_query(query)
        if not result.get("success"):
            return result

        return {
            "table_name": table_name,
            "schema": schema,
            "sample": result.get("rows", []),
            "sample_size": len(result.get("rows", [])),
        }

    async def get_column_stats(
        self,
        table_name: str,
        column_name: str,
        schema: str = "public",
    ) -> dict[str, Any]:
        """Get statistics for a specific column."""
        # Get basic stats
        query = f'''
        SELECT
            COUNT(*) as total_count,
            COUNT("{column_name}") as non_null_count,
            COUNT(DISTINCT "{column_name}") as distinct_count
        FROM "{schema}"."{table_name}"
        '''
        result = await self.execute_query(query)
        if not result.get("success"):
            return result

        row = result.get("rows", [{}])[0]
        total = row.get("total_count", 0)
        non_null = row.get("non_null_count", 0)

        return {
            "table_name": table_name,
            "column_name": column_name,
            "total_count": total,
            "non_null_count": non_null,
            "null_count": total - non_null,
            "null_rate": (total - non_null) / max(total, 1),
            "distinct_count": row.get("distinct_count", 0),
            "cardinality": row.get("distinct_count", 0) / max(non_null, 1),
        }


# Tool definitions for Claude Agent SDK
DATABASE_TOOLS = [
    {
        "name": "execute_query",
        "description": "Execute a SQL query on the database. Use with caution - only SELECT queries recommended for safety.",
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "SQL query to execute",
                },
                "params": {
                    "type": "array",
                    "description": "Query parameters (for parameterized queries)",
                    "items": {"type": "string"},
                },
                "fetch": {
                    "type": "boolean",
                    "description": "Whether to fetch and return results",
                    "default": True,
                },
            },
            "required": ["query"],
        },
    },
    {
        "name": "get_table_info",
        "description": "Get column information for a database table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Name of the table"},
                "schema": {"type": "string", "description": "Schema name", "default": "public"},
            },
            "required": ["table_name"],
        },
    },
    {
        "name": "get_table_stats",
        "description": "Get statistics (row count, size) for a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Name of the table"},
                "schema": {"type": "string", "description": "Schema name", "default": "public"},
            },
            "required": ["table_name"],
        },
    },
    {
        "name": "list_tables",
        "description": "List all tables in a database schema.",
        "input_schema": {
            "type": "object",
            "properties": {
                "schema": {"type": "string", "description": "Schema name", "default": "public"},
            },
        },
    },
    {
        "name": "sample_data",
        "description": "Get a sample of rows from a table.",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Name of the table"},
                "schema": {"type": "string", "description": "Schema name", "default": "public"},
                "limit": {"type": "integer", "description": "Number of rows to sample", "default": 10},
            },
            "required": ["table_name"],
        },
    },
    {
        "name": "get_column_stats",
        "description": "Get statistics for a specific column (null rate, cardinality, distinct count).",
        "input_schema": {
            "type": "object",
            "properties": {
                "table_name": {"type": "string", "description": "Name of the table"},
                "column_name": {"type": "string", "description": "Name of the column"},
                "schema": {"type": "string", "description": "Schema name", "default": "public"},
            },
            "required": ["table_name", "column_name"],
        },
    },
]


__all__ = ["DatabaseTools", "DATABASE_TOOLS"]
