"""
Database connector for TiDB using SQLAlchemy.
Provides functionality to connect to TiDB and fetch plan data.
"""

from pathlib import Path
from typing import Dict, List

import pandas as pd
from sqlalchemy import create_engine, text


def load_env_file(env_path: str = None) -> Dict[str, str]:
    """
    Load environment variables from .env file.

    Args:
        env_path: Path to .env file (default: look for .env in current directory)

    Returns:
        Dictionary of environment variables
    """
    if env_path is None:
        # Look for .env file in the same directory as this script
        current_dir = Path(__file__).parent
        env_path = current_dir / ".env"

    env_vars = {}

    if Path(env_path).exists():
        with open(env_path, "r") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    # Remove quotes if present
                    value = value.strip("\"'")
                    env_vars[key] = value

    return env_vars


class TiDBConnector:
    """SQLAlchemy connector for TiDB database."""

    def __init__(
        self,
        host: str = None,
        port: int = None,
        user: str = None,
        password: str = None,
        database: str = None,
        env_file: str = None,
    ):
        """
        Initialize TiDB connector.

        Args:
            host: Database host (if None, will load from .env)
            port: Database port (if None, will load from .env)
            user: Database username (if None, will load from .env)
            password: Database password (if None, will load from .env)
            database: Database name (if None, will load from .env)
            env_file: Path to .env file (optional)
        """
        # Load from .env file if parameters not provided
        if any(param is None for param in [host, port, user, password, database]):
            env_vars = load_env_file(env_file)

            self.host = host or env_vars.get("DB_URL", "localhost")
            self.port = int(port or env_vars.get("DB_PORT", "4000"))
            self.user = user or env_vars.get("DB_USERNAME", "root")
            self.password = password or env_vars.get("DB_PASSWORD", "")
            self.database = database or env_vars.get("DEFAULT_DB", "test")
        else:
            self.host = host
            self.port = port
            self.user = user
            self.password = password
            self.database = database

        self.engine = None
        self._connection_string = None

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

        Returns:
            True if connection successful, False otherwise
        """
        try:
            # Build connection string for TiDB (MySQL protocol)
            self._connection_string = (
                f"mysql+pymysql://{self.user}:{self.password}@"
                f"{self.host}:{self.port}/{self.database}"
            )

            # Create engine with connection pooling
            self.engine = create_engine(
                self._connection_string,
                pool_size=5,
                max_overflow=10,
                pool_pre_ping=True,
                pool_recycle=3600,
            )

            # Test connection
            with self.engine.connect() as conn:
                result = conn.execute(text("SELECT 1 as test"))
                result.fetchone()

            print(f"✓ Successfully connected to TiDB at {self.host}:{self.port}")
            return True

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

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

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

        try:
            with self.engine.connect() as conn:
                # Test basic connectivity
                version_result = conn.execute(text("SELECT VERSION() as version"))
                version = version_result.fetchone()[0]

                # Get current database
                db_result = conn.execute(text("SELECT DATABASE() as db"))
                current_db = db_result.fetchone()[0]

                # Check if statements summary table exists
                table_check = conn.execute(
                    text(
                        """
                    SELECT COUNT(*) as count
                    FROM information_schema.tables
                    WHERE table_schema = 'INFORMATION_SCHEMA'
                    AND table_name = 'CLUSTER_STATEMENTS_SUMMARY_HISTORY'
                """
                    )
                )
                has_statements_table = table_check.fetchone()[0] > 0

                return {
                    "connected": True,
                    "version": version,
                    "current_database": current_db,
                    "has_statements_table": has_statements_table,
                    "connection_string": self._connection_string.replace(
                        self.password, "***"
                    ),
                }

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

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

        Args:
            limit: Maximum number of records to fetch

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

        query = text(
            """
            SELECT
                STMT_TYPE,
                TABLE_NAMES,
                INDEX_NAMES,
                QUERY_SAMPLE_TEXT,
                PLAN,
                PLAN_HINT
            FROM INFORMATION_SCHEMA.CLUSTER_STATEMENTS_SUMMARY_HISTORY
            WHERE PLAN IS NOT NULL
            AND PLAN != ''
            ORDER BY SUMMARY_BEGIN_TIME DESC
            LIMIT :limit
        """
        )

        try:
            with self.engine.connect() as conn:
                result = conn.execute(query, {"limit": limit})
                records = []

                for row in result:
                    records.append(
                        {
                            "stmt_type": row[0],
                            "table_names": row[1],
                            "index_names": row[2],
                            "sql": row[3],
                            "plan_text": row[4],
                            "plan_hint": row[5],
                            "version": "v4",
                        }
                    )

                return records

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

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

        Args:
            limit: Maximum number of records to fetch

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

    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.engine:
            raise RuntimeError("Database not connected")

        try:
            with self.engine.connect() as conn:
                query = text(query_str)
                result = conn.execute(query, params or {})

                # Get column names
                columns = list(result.keys())

                # Convert to list of dictionaries
                records = []
                for row in result:
                    record = {}
                    for i, col in enumerate(columns):
                        record[col] = row[i]
                    records.append(record)

                return records

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

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


def test_env_connection():
    """Test connection to TiDB using .env configuration."""
    print("Testing TiDB connection using .env configuration...")

    # Create connector using .env file
    connector = TiDBConnector()

    print(f"Connecting to: {connector.host}:{connector.port}")
    print(f"Database: {connector.database}")
    print(f"User: {connector.user}")

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

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

            if statements:
                print("\nFirst statement:")
                stmt = statements[0]
                print(f"  Type: {stmt['stmt_type']}")
                print(f"  SQL: {stmt['sql'][:100]}...")
                print(f"  Plan length: {len(stmt['plan_text'])} chars")

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

        finally:
            connector.close()

    else:
        print("Connection failed!")


# Convenience functions for quick usage
def execute_query(
    query_str: str, params: Dict = None, database: 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)
        auto_close: Automatically close connection after query (default: True)

    Returns:
        List of dictionaries representing query results
    """
    connector = TiDBConnector(database=database)

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

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


def fetch_workload_plans(limit: int = 10) -> List[Dict]:
    """
    Convenience function to fetch plans from workload backup table.

    Args:
        limit: Number of plans to fetch

    Returns:
        List of plan dictionaries
    """
    query = """
        SELECT STMT_TYPE, QUERY_SAMPLE_TEXT, PLAN, LENGTH(PLAN) as plan_length
        FROM backup.workload_statements_history_backup
        WHERE PLAN IS NOT NULL AND PLAN != ''
        AND LENGTH(PLAN) > 100
        ORDER BY LENGTH(PLAN) DESC
        LIMIT :limit
    """

    results = execute_query(query, {"limit": limit})

    # Rename columns to match expected format
    plans = []
    for row in results:
        plans.append(
            {
                "stmt_type": row["STMT_TYPE"],
                "sql": row["QUERY_SAMPLE_TEXT"],
                "plan_text": row["PLAN"],
                "plan_length": row["plan_length"],
            }
        )

    return plans


if __name__ == "__main__":
    test_env_connection()