"""Context object passed to doctor diagnostic rules providing shared connections."""
from __future__ import annotations

from typing import Any, Dict, List
import os

from sqlalchemy import create_engine, MetaData, Table, select, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import SQLAlchemyError


class DoctorContext:
    """Provides Database and Prometheus helper connections for rules."""
    def __init__(self, engine: Engine, prometheus_url: str = None, table_name: str = None):
        self.engine = engine
        self.prometheus_url = prometheus_url or 'http://localhost:9090'
        self.table_name = table_name
        self.metadata = MetaData()
        self._table = None
        
        # Only try to reflect the table if table_name is provided
        if table_name:
            try:
                self._table = Table(table_name, self.metadata, autoload_with=engine)
            except Exception as e:
                print(f"Warning: Could not reflect table '{table_name}': {e}")
                print("You can still use raw SQL queries with execute_sql()")

    @classmethod
    def create(cls) -> 'DoctorContext':
        """Factory that creates context from environment variables."""
        # Get database URL from environment or use SQLite in-memory as default
        # db_url = os.getenv('DOCTOR_DB_URL', 'sqlite:///:memory:')
        db_url = os.getenv('DOCTOR_DB_URL', 'mysql+pymysql://root:changeit@localhost/o11y')
        prometheus_url = os.getenv('PROMETHEUS_URL')
        table_name = os.getenv('DOCTOR_TABLE_NAME')

        # Create database engine with connection pooling
        engine = create_engine(db_url, pool_pre_ping=True)
        return cls(engine, prometheus_url, table_name)

    def execute_sql(self, sql: str, params: dict = None) -> List[Dict[str, Any]]:
        """Execute raw SQL query and return results as a list of dictionaries.
        
        Args:
            sql: SQL query string
            params: Optional parameters for the query
            
        Returns:
            List of dictionaries representing the query results
        """
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text(sql), params or {})
                if result.returns_rows:
                    columns = list(result.keys())
                    return [dict(zip(columns, row)) for row in result]
                return []
        except SQLAlchemyError as e:
            print(f"Error executing SQL: {e}")
            raise

    def query(self, *args, **kwargs) -> List[Dict[str, Any]]:
        """Execute a query against the configured table (if any).
        
        Args:
            *args: Positional arguments passed to select()
            **kwargs: Keyword arguments for filtering (column=value)
            
        Returns:
            List of dictionaries representing the query results
            
        Raises:
            ValueError: If no table is configured
        """
        if not self._table:
            raise ValueError("No table configured. Use execute_sql() for raw SQL queries.")
            
        stmt = select(*args) if args else select(self._table)
        if kwargs:
            stmt = stmt.filter_by(**kwargs)
            
        with self.engine.connect() as conn:
            result = conn.execute(stmt)
            return [dict(row._mapping) for row in result]

    def prom(self, expr: str):
        """Query Prometheus (stub implementation)."""
        # TODO: Implement or remove Prometheus integration
        print(f"[DEBUG] Would query Prometheus at {self.prometheus_url} with: {expr}")
        return []