"""
SQLite connector for SimCluster evaluation system.
Provides functionality to connect to SQLite database and fetch plan data.
"""

import sqlite3
from pathlib import Path
from typing import Dict, List

import pandas as pd


class SQLiteConnector:
    """SQLite connector for SimCluster database."""

    def __init__(self, db_path: str = None):
        """
        Initialize SQLite connector.

        Args:
            db_path: Path to SQLite database file (default: simcluster.db in current directory)
        """
        if db_path is None:
            # Default to simcluster.db in current directory
            db_path = Path.cwd() / "simcluster.db"

        self.db_path = Path(db_path)
        self.connection = None

    def connect(self) -> bool:
        """
        Establish connection to SQLite database.

        Returns:
            True if connection successful, False otherwise
        """
        try:
            if not self.db_path.exists():
                print(f"✗ SQLite database not found: {self.db_path}")
                return False

            # Create connection with row factory for dict-like access
            self.connection = sqlite3.connect(str(self.db_path))
            self.connection.row_factory = sqlite3.Row

            # Test connection
            cursor = self.connection.cursor()
            cursor.execute("SELECT 1 as test")
            cursor.fetchone()

            print(f"✓ Successfully connected to SQLite: {self.db_path}")
            return True

        except Exception as e:
            print(f"✗ Failed to connect to SQLite: {e}")
            return False

    def test_connection(self) -> Dict:
        """
        Test database connection and get basic info.

        Returns:
            Dictionary with connection test results
        """
        if not self.connection:
            return {"connected": False, "error": "No connection established"}

        try:
            cursor = self.connection.cursor()

            # Get SQLite version
            cursor.execute("SELECT sqlite_version() as version")
            version = cursor.fetchone()[0]

            # Check if statements summary table exists
            cursor.execute(
                """
                SELECT COUNT(*) as count
                FROM sqlite_master
                WHERE type='table'
                AND name='cluster_statements_summary_history'
            """
            )
            has_statements_table = cursor.fetchone()[0] > 0

            # Get record count if table exists
            record_count = 0
            if has_statements_table:
                cursor.execute(
                    "SELECT COUNT(*) FROM cluster_statements_summary_history"
                )
                record_count = cursor.fetchone()[0]

            # Get schema distribution
            schema_distribution = {}
            if has_statements_table:
                cursor.execute(
                    """
                    SELECT schema_name, COUNT(*) as count
                    FROM cluster_statements_summary_history
                    GROUP BY schema_name
                """
                )
                schema_distribution = {row[0]: row[1] for row in cursor.fetchall()}

            return {
                "connected": True,
                "sqlite_version": version,
                "database_path": str(self.db_path),
                "has_statements_table": has_statements_table,
                "record_count": record_count,
                "schema_distribution": schema_distribution,
            }

        except Exception as e:
            return {"connected": False, "error": str(e)}

    def execute_query(self, query_str: str, params: Dict = None) -> List[Dict]:
        """
        Execute a SQL query and return results as list of dictionaries.

        Args:
            query_str: SQL query string
            params: Query parameters (optional)

        Returns:
            List of dictionaries representing query results
        """
        if not self.connection:
            raise RuntimeError("Database not connected")

        try:
            cursor = self.connection.cursor()

            # Convert named parameters to SQLite format if needed
            if params:
                # SQLite uses :param_name format
                for key, value in params.items():
                    if f":{key}" not in query_str and f"{{{key}}}" in query_str:
                        query_str = query_str.replace(f"{{{key}}}", f":{key}")

            cursor.execute(query_str, params or {})

            # Convert rows to list of dictionaries
            records = []
            for row in cursor.fetchall():
                # sqlite3.Row supports dict-like access
                record = dict(row)
                records.append(record)

            return records

        except Exception as e:
            raise RuntimeError(f"SQLite query failed: {e}")

    def fetch_statements_summary(
        self, schema_name: str = "test", limit: int = 100
    ) -> List[Dict]:
        """
        Fetch statements from cluster_statements_summary_history table.

        Args:
            schema_name: Schema name to filter by (default: test)
            limit: Maximum number of records to fetch

        Returns:
            List of statement records
        """
        if not self.connection:
            raise RuntimeError("Database not connected")

        query = """
            SELECT
                stmt_type,
                table_names,
                index_names,
                query_sample_text,
                plan,
                plan_hint,
                digest_text,
                exec_count,
                avg_latency,
                schema_name
            FROM cluster_statements_summary_history
            WHERE plan IS NOT NULL
            AND plan != ''
            AND schema_name = :schema_name
            ORDER BY summary_begin_time DESC
            LIMIT :limit
        """

        try:
            records = self.execute_query(
                query, {"schema_name": schema_name, "limit": limit}
            )

            # Rename columns to match expected format
            formatted_records = []
            for record in records:
                formatted_records.append(
                    {
                        "stmt_type": record["stmt_type"],
                        "table_names": record["table_names"],
                        "index_names": record["index_names"],
                        "sql": record["query_sample_text"],
                        "plan_text": record["plan"],
                        "plan_hint": record["plan_hint"],
                        "digest_text": record["digest_text"],
                        "exec_count": record["exec_count"],
                        "avg_latency": record["avg_latency"],
                        "schema_name": record["schema_name"],
                        "version": "v4",
                    }
                )

            return formatted_records

        except Exception as e:
            raise RuntimeError(f"Failed to fetch statements: {e}")

    def fetch_statements_summary_df(
        self, schema_name: str = "test", limit: int = 100
    ) -> pd.DataFrame:
        """
        Fetch statements as pandas DataFrame.

        Args:
            schema_name: Schema name to filter by
            limit: Maximum number of records to fetch

        Returns:
            DataFrame with statement records
        """
        records = self.fetch_statements_summary(schema_name, limit)
        return pd.DataFrame(records)

    def get_schema_stats(self) -> Dict[str, int]:
        """
        Get statistics about schemas in the database.

        Returns:
            Dictionary with schema names and record counts
        """
        if not self.connection:
            raise RuntimeError("Database not connected")

        query = """
            SELECT
                schema_name,
                COUNT(*) as total_count,
                COUNT(CASE WHEN plan IS NOT NULL AND plan != '' THEN 1 END) as with_plan_count
            FROM cluster_statements_summary_history
            GROUP BY schema_name
        """

        try:
            records = self.execute_query(query)
            stats = {}
            for record in records:
                stats[record["schema_name"]] = {
                    "total_count": record["total_count"],
                    "with_plan_count": record["with_plan_count"],
                }
            return stats

        except Exception as e:
            raise RuntimeError(f"Failed to get schema stats: {e}")

    def close(self):
        """Close database connection."""
        if self.connection:
            self.connection.close()
            self.connection = None
            print("SQLite connection closed")


def test_sqlite_connection(db_path: str = None):
    """Test connection to SQLite database."""
    print("Testing SQLite connection...")

    # Create connector
    connector = SQLiteConnector(db_path)

    print(f"Database path: {connector.db_path}")

    # Test connection
    if connector.connect():
        # Get connection info
        info = connector.test_connection()
        print("\nConnection Info:")
        for key, value in info.items():
            if isinstance(value, dict):
                print(f"  {key}:")
                for sub_key, sub_value in value.items():
                    print(f"    {sub_key}: {sub_value}")
            else:
                print(f"  {key}: {value}")

        # Get schema statistics
        try:
            print("\nSchema Statistics:")
            stats = connector.get_schema_stats()
            for schema, schema_stats in stats.items():
                print(
                    f"  {schema}: {schema_stats['with_plan_count']}/{schema_stats['total_count']} records with plans"
                )

        except Exception as e:
            print(f"Warning: Could not get schema stats: {e}")

        # Try to fetch some data
        try:
            print("\nFetching sample statements...")
            statements = connector.fetch_statements_summary(limit=3)
            print(f"Found {len(statements)} statements")

            if statements:
                print("\nFirst statement:")
                stmt = statements[0]
                print(f"  Schema: {stmt['schema_name']}")
                print(f"  Type: {stmt['stmt_type']}")
                print(f"  SQL: {stmt['sql'][:100]}...")
                print(f"  Plan length: {len(stmt['plan_text'])} chars")
                print(f"  Exec count: {stmt['exec_count']}")
                print(f"  Avg latency: {stmt['avg_latency']}")

        except Exception as e:
            print(f"Warning: Could not fetch statements: {e}")

        finally:
            connector.close()

    else:
        print("Connection failed!")


# Convenience function for quick usage
def execute_sqlite_query(
    query_str: str, params: Dict = None, db_path: str = None, auto_close: bool = True
) -> List[Dict]:
    """
    Execute a query with automatic connection management.

    Args:
        query_str: SQL query string
        params: Query parameters (optional)
        db_path: Database path (optional)
        auto_close: Automatically close connection after query (default: True)

    Returns:
        List of dictionaries representing query results
    """
    connector = SQLiteConnector(db_path)

    if not connector.connect():
        raise RuntimeError("Failed to connect to SQLite database")

    try:
        return connector.execute_query(query_str, params)
    finally:
        if auto_close:
            connector.close()


if __name__ == "__main__":
    test_sqlite_connection()