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

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()


def get_sql_engine(
    db_host=None,
    db_user=None,
    db_pass=None,
    db_name=None,
    db_port=5432,
    instance_connection_name=None,
) -> sqlalchemy.engine.base.Engine:
    """Initializes a connection pool for a Cloud SQL instance.

    If instance_connection_name is provided (or found in INSTANCE_CONNECTION_NAME env var),
    it uses the Cloud SQL Python Connector with IAM authentication.
    Otherwise, it falls back to a TCP connection.

    Args:
        db_host: Database host address (e.g., '127.0.0.1')
        db_user: Database username (for IAM, defaults to detecting the service account email)
        db_pass: Database password (not used for IAM)
        db_name: Database name (defaults to DB_NAME env var)
        db_port: Database port (default: 5432)
        instance_connection_name: Cloud SQL instance connection name (PROJECT:REGION:INSTANCE)

    Returns:
        SQLAlchemy Engine instance
    """
    # Try to get defaults from environment if not provided
    instance_connection_name = instance_connection_name or os.getenv(
        "INSTANCE_CONNECTION_NAME"
    )
    db_name = db_name or os.getenv("DB_NAME")

    if instance_connection_name:
        # If db_user is not provided, try to detect the service account email
        if not db_user:
            try:
                credentials, _ = google.auth.default()
                if hasattr(credentials, "service_account_email"):
                    # For Postgres IAM, the user is the service account email without .gserviceaccount.com
                    sa_email = credentials.service_account_email
                    db_user = sa_email.replace(".gserviceaccount.com", "")
                    logger.info(
                        f"Auto-detected service account for IAM auth: {db_user}"
                    )
                elif hasattr(credentials, "signer_email"):
                    # Fallback for some credential types
                    db_user = credentials.signer_email.replace(
                        ".gserviceaccount.com", ""
                    )
                    logger.info(f"Auto-detected service account from signer: {db_user}")
            except Exception as e:
                logger.warning(f"Could not auto-detect service account email: {e}")

        logger.info(
            f"Connecting to Cloud SQL instance '{instance_connection_name}' using IAM auth as user '{db_user}'."
        )

        # initialize Connector object
        connector = Connector()

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

        pool = sqlalchemy.create_engine(
            "postgresql+pg8000://",
            creator=getconn,
        )
        logger.info("Successfully created Cloud SQL connector connection pool.")
        return pool

    logger.info(f"Connecting to database '{db_name}' at {db_host}:{db_port}")

    try:
        pool = sqlalchemy.create_engine(
            # Equivalent URL:
            # postgresql+pg8000://<db_user>:<db_pass>@<db_host>:<db_port>/<db_name>
            sqlalchemy.engine.url.URL.create(
                drivername="postgresql+pg8000",
                username=db_user,
                password=db_pass,
                host=db_host,
                port=db_port,
                database=db_name,
            ),
        )
        logger.info("Successfully created TCP database connection pool.")
        return pool
    except Exception as e:
        logger.exception(f"Error connecting to database: {str(e)}")
        raise e


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

    Args:
        engine: SQLAlchemy Engine instance

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


def execute_query(engine, 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 = get_session(engine)
    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(engine, base_class=Base):
    """Creates all tables defined in the SQLAlchemy Base metadata.

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


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

    Args:
        engine: SQLAlchemy Engine instance
        model_instance: SQLAlchemy model instance to save

    Returns:
        The saved model instance with refreshed data
    """
    session = get_session(engine)
    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(engine, model_class):
    """Retrieves all records for a given model class.

    Args:
        engine: SQLAlchemy Engine instance
        model_class: SQLAlchemy model class

    Returns:
        List of all records for the model
    """
    session = get_session(engine)
    try:
        logger.info(f"Retrieving all records for {model_class.__name__}")
        results = session.query(model_class).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(engine, model_class, record_id):
    """Retrieves a single record by its primary key ID.

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

    Returns:
        Model instance if found, None otherwise
    """
    session = get_session(engine)
    try:
        logger.info(f"Retrieving {model_class.__name__} with ID: {record_id}")
        result = session.query(model_class).get(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(engine, model_class, **filters):
    """Retrieves records matching the provided filter criteria.

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

    Returns:
        List of matching records
    """
    session = get_session(engine)
    try:
        logger.info(f"Filtering {model_class.__name__} with criteria: {filters}")
        results = session.query(model_class).filter_by(**filters).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(engine, model_class, **filters):
    """Retrieves a single record matching the provided filter criteria.

    Args:
        engine: SQLAlchemy Engine instance
        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 = get_session(engine)
    try:
        logger.info(
            f"Retrieving single {model_class.__name__} with criteria: {filters}"
        )
        result = session.query(model_class).filter_by(**filters).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(engine, model_instance):
    """Updates an existing model instance in the database.

    Args:
        engine: SQLAlchemy Engine instance
        model_instance: SQLAlchemy model instance with updated values

    Returns:
        The updated model instance with refreshed data
    """
    session = get_session(engine)
    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(engine, model_instance):
    """Deletes a model instance from the database.

    Args:
        engine: SQLAlchemy Engine instance
        model_instance: SQLAlchemy model instance to delete
    """
    session = get_session(engine)
    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(engine):
    """Disposes the connection pool for the provided engine.

    Args:
        engine: SQLAlchemy Engine instance to dispose
    """
    try:
        logger.info("Disposing database connection pool...")
        engine.dispose()
        logger.info("Successfully disposed connection pool.")
    except Exception as e:
        logger.exception(f"Error disposing connection pool: {str(e)}")
        raise e
