import os
import logging
import sqlalchemy
from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base
from sqlalchemy import text
from sqlalchemy.engine.base import Engine

import google.auth
from google.cloud.sql.connector import Connector, IPTypes

# Configure the logger
logger = logging.getLogger("uvicorn")
logger.setLevel(logging.INFO)

Base = declarative_base()

import os
import logging
import sqlalchemy
from sqlalchemy.engine import Engine
from google.cloud.sql.connector import Connector, IPTypes
import google.auth
from google.auth.transport.requests import Request

logger = logging.getLogger(__name__)


def _detect_service_account_email() -> str:
    credentials, _ = google.auth.default()
    credentials.refresh(Request())

    email = getattr(credentials, "service_account_email", None) or getattr(
        credentials, "signer_email", None
    )

    if not email:
        raise RuntimeError("Could not determine service account email for IAM login")

    return email


def get_sql_engine(
    db_host=None,
    db_user=None,
    db_pass=None,
    db_name=None,
    db_port=5432,
    instance_connection_name=None,
    ip_type=IPTypes.PSC,
) -> Engine:
    """
    Create SQLAlchemy engine.

    Uses Cloud SQL Connector + IAM if instance_connection_name is provided.
    Otherwise falls back to TCP connection.
    """
    if instance_connection_name:
        if not db_user:
            db_user = _detect_service_account_email()
        
        logger.info(
            f"Connecting via Cloud SQL Connector (ip_type={ip_type.name}) "
            f"to instance '{instance_connection_name}' as '{db_user}'"
        )

        connector = Connector()

        def getconn():
            return connector.connect(
                instance_connection_name,
                "pg8000",
                user=db_user,
                db=db_name,
                enable_iam_auth=True,
                ip_type=ip_type,
            )

        engine = sqlalchemy.create_engine(
            "postgresql+pg8000://",
            creator=getconn,
            pool_pre_ping=True,
        )

        return engine

    # ---------- TCP fallback ----------

    if not all([db_host, db_user, db_pass, db_name]):
        raise ValueError(
            "TCP connection requires db_host, db_user, db_pass, db_name"
        )

    logger.info(f"Connecting via TCP to {db_host}:{db_port}/{db_name}")

    return sqlalchemy.create_engine(
        sqlalchemy.engine.URL.create(
            drivername="postgresql+pg8000",
            username=db_user,
            password=db_pass,
            host=db_host,
            port=db_port,
            database=db_name,
        ),
        pool_pre_ping=True,
    )

class GCPCloudSQL:
    """
    Utility class for managing a connection to a Cloud SQL instance.
    For creation you must supply an engine.
    see get_sql_engine defined above for creating an engine
    """
    engine: Engine

    def __init__(self, engine: Engine):
        self.engine = engine

    def get_engine(self):
        return self.engine

    def get_session(self):
        """Creates and returns a new session for the provided engine.

        Args:
            engine: SQLAlchemy Engine instance

        Returns:
            SQLAlchemy Session instance
        """
        Session = sessionmaker(bind=self.engine)
        return Session()


    def execute_query(self, query_str: str, params: dict = None):
        """Executes a raw SQL query.

        Args:
            engine: SQLAlchemy Engine instance
            query_str: SQL query string
            params: Dictionary of query parameters (optional)

        Returns:
            Query result
        """
        session = self.get_session()
        try:
            logger.info(f"Executing query: {query_str[:100]}...")
            result = session.execute(text(query_str), params or {})
            session.commit()
            logger.info("Query executed successfully.")
            return result
        except Exception as e:
            session.rollback()
            logger.exception(f"Error executing query: {str(e)}")
            raise e
        finally:
            session.close()


    def create_tables(self, base_class=Base):
        """Creates all tables defined in the SQLAlchemy Base metadata.

        Args:
            base_class: SQLAlchemy declarative base class (default: Base)
        """
        try:
            logger.info("Creating database tables...")
            base_class.metadata.create_all(self.engine)
            logger.info("Successfully created database tables.")
        except Exception as e:
            logger.exception(f"Error creating tables: {str(e)}")
            raise e


    def save_model(self, model_instance):
        """Saves a single model instance to the database.

        Args:
            model_instance: SQLAlchemy model instance to save

        Returns:
            The saved model instance with refreshed data
        """
        session = self.get_session()
        try:
            logger.info(f"Saving model instance: {type(model_instance).__name__}")
            session.add(model_instance)
            session.commit()
            session.refresh(model_instance)
            logger.info("Successfully saved model instance.")
            return model_instance
        except Exception as e:
            session.rollback()
            logger.exception(f"Error saving model: {str(e)}")
            raise e
        finally:
            session.close()


    def get_all(self, model_class):
        """Retrieves all records for a given model class.

        Args:
            model_class: SQLAlchemy model class

        Returns:
            List of all records for the model
        """
        session = self.get_session()
        try:
            logger.info(f"Retrieving all records for {model_class.__name__}")
            statement = sqlalchemy.select(model_class)
            results = session.execute(statement).scalars().all()
            logger.info(f"Retrieved {len(results)} records.")
            return results
        except Exception as e:
            logger.exception(f"Error retrieving records: {str(e)}")
            raise e
        finally:
            session.close()


    def get_by_id(self, model_class, record_id):
        """Retrieves a single record by its primary key ID.

        Args:
            model_class: SQLAlchemy model class
            record_id: Primary key value to search for

        Returns:
            Model instance if found, None otherwise
        """
        session = self.get_session()
        try:
            logger.info(f"Retrieving {model_class.__name__} with ID: {record_id}")
            result = session.get(model_class, record_id)
            if result:
                logger.info(f"Found record with ID: {record_id}")
            else:
                logger.info(f"No record found with ID: {record_id}")
            return result
        except Exception as e:
            logger.exception(f"Error retrieving record by ID: {str(e)}")
            raise e
        finally:
            session.close()


    def get_by_filter(self, model_class, **filters):
        """Retrieves records matching the provided filter criteria.

        Args:
            model_class: SQLAlchemy model class
            **filters: Keyword arguments for filtering (e.g., name="John", age=30)

        Returns:
            List of matching records
        """
        session = self.get_session()
        try:
            logger.info(f"Filtering {model_class.__name__} with criteria: {filters}")
            statement = sqlalchemy.select(model_class).filter_by(**filters)
            results = session.execute(statement).scalars().all()
            logger.info(f"Found {len(results)} matching records.")
            return results
        except Exception as e:
            logger.exception(f"Error filtering records: {str(e)}")
            raise e
        finally:
            session.close()


    def get_one(self, model_class, **filters):
        """Retrieves a single record matching the provided filter criteria.

        Args:
            model_class: SQLAlchemy model class
            **filters: Keyword arguments for filtering (e.g., email="user@example.com")

        Returns:
            First matching model instance if found, None otherwise
        """
        session = self.get_session()
        try:
            logger.info(
                f"Retrieving single {model_class.__name__} with criteria: {filters}"
            )
            statement = sqlalchemy.select(model_class).filter_by(**filters)
            result = session.execute(statement).scalars().first()
            if result:
                logger.info(f"Found matching record.")
            else:
                logger.info(f"No matching record found.")
            return result
        except Exception as e:
            logger.exception(f"Error retrieving single record: {str(e)}")
            raise e
        finally:
            session.close()


    def update_model(self, model_instance):
        """Updates an existing model instance in the database.

        Args:
            model_instance: SQLAlchemy model instance with updated values

        Returns:
            The updated model instance with refreshed data
        """
        session = self.get_session()
        try:
            logger.info(f"Updating model instance: {type(model_instance).__name__}")
            session.merge(model_instance)
            session.commit()
            logger.info("Successfully updated model instance.")
            return model_instance
        except Exception as e:
            session.rollback()
            logger.exception(f"Error updating model: {str(e)}")
            raise e
        finally:
            session.close()


    def delete_model(self, model_instance):
        """Deletes a model instance from the database.

        Args:
            model_instance: SQLAlchemy model instance to delete
        """
        session = self.get_session()
        try:
            logger.info(f"Deleting model instance: {type(model_instance).__name__}")
            session.delete(model_instance)
            session.commit()
            logger.info("Successfully deleted model instance.")
        except Exception as e:
            session.rollback()
            logger.exception(f"Error deleting model: {str(e)}")
            raise e
        finally:
            session.close()


    def cleanup_connector(self):
        """Disposes the connection pool for the provided engine."""
        try:
            logger.info("Disposing database connection pool...")
            self.engine.dispose()
            logger.info("Successfully disposed connection pool.")
        except Exception as e:
            logger.exception(f"Error disposing connection pool: {str(e)}")
            raise e
