import logging
import time
from typing import Any, Dict, Optional, Tuple, TypeVar

from sqlalchemy import create_engine, Engine, Result, text
from sqlalchemy.orm import scoped_session, Session, sessionmaker

from tests import config
from tests.utils.secrets_manager import SecretsManager

log = logging.getLogger(__name__)
# Define a generic type for the result row
T = TypeVar('T', bound=Tuple[Any, ...])  # Represents a row structure


# Initialize engine only once
def get_engine() -> Engine:
    """
    Create and return a SQLAlchemy engine
    """
    db_name = config.DB_NAME
    credentials = SecretsManager().get_db_credentials()['db'][db_name]

    user = credentials['user']
    password = credentials['password']
    host = credentials['host']
    port = credentials['port']

    sqlalchemy_database_uri = (
        f'mysql+pymysql://{user}:{password}@{host}:{port}/{db_name}'
    )

    return create_engine(
        sqlalchemy_database_uri,
        pool_pre_ping=True,
        pool_recycle=3600,
        echo=False,  # Set to True only for debugging
        ## MySQL uses REPEATABLE READ by default, and causes stale reads. Use READ COMMITTED to see changes from other transactions.
        isolation_level='READ COMMITTED',
    )


# Create and cache engine and scoped session factory only once
_engine: Engine = get_engine()

_SessionFactory: scoped_session[Session] = scoped_session(
    sessionmaker(
        bind=_engine,
        autocommit=False,
        autoflush=False,
        expire_on_commit=True,
    )
)


def create_session() -> scoped_session[Session]:
    """
    Return a scoped_session factory.
    """
    return _SessionFactory


def construct_where_clause(conditions: Dict[str, Any]) -> str:
    return ' AND '.join([f'{key}=:{key}' for key in conditions])


def construct_set_clause(values: Dict[str, Any]) -> str:
    return ', '.join(f'{key} = :{key}' for key in values)


def execute_query(db_session: Session, query: str, params: Dict[str, Any]) -> Result[T]:
    """Executes the given SQL query and returns the raw result."""
    return db_session.execute(text(query), params)


def process_result(result: Result[T]) -> Optional[Dict[str, Any]]:
    """Processes the query result and returns a dictionary if data exists."""
    row = result.mappings().fetchone()

    if row:
        return dict(row)

    log.info('Entity not found with given conditions')
    return None


def get_max_value(
    db_session: Session, table_name: str, column_name: str, conditions: Dict[str, Any]
) -> Optional[Any]:
    """Fetches the MAX value of a specific column from the given table."""
    where_clause = construct_where_clause(conditions)
    """https://github.com/theorchard/collab/blob/b335eb5b1543ebaf9a08be9c23d34c985929b6d0/abacus/refresh/data-refresh/uk_wht_tax_data_for_testing.sql#L66-L68 """
    query = f'SELECT MAX({column_name}) FROM {table_name} WHERE {where_clause}'

    result: Result[Tuple[Optional[Any]]] = execute_query(db_session, query, conditions)
    return result.scalar()


def get_entity(
    db_session: Session, table_name: str, conditions: Dict[str, Any]
) -> Optional[Dict[str, Any]]:
    """Fetches an entity (all columns) from the specified table using conditions."""
    where_clause = construct_where_clause(conditions)
    query = f'SELECT * FROM {table_name} WHERE {where_clause}'

    result: Result[Tuple[Optional[Any]]] = execute_query(db_session, query, conditions)
    return process_result(result)


def update_entity(
    db_session: Session,
    table_name: str,
    conditions: Dict[str, Any],
    values: Dict[str, Any],
) -> None:
    """
    Updates rows in the specified table matching conditions with provided values.
    """
    set_clause = construct_set_clause(values)
    where_clause = construct_where_clause(conditions)
    update_query = f'UPDATE {table_name} SET {set_clause} WHERE {where_clause}'

    # Merge parameters for SET and WHERE clauses
    params = {**values, **conditions}
    execute_query(db_session, update_query, params)

    return None


def delete_entity_by_id(
    db_session: Session, table_name: str, row_name: str, entity_id: int
) -> None:
    """Deletes an entity from the specified table by ID using a database session."""
    delete_query = f'DELETE FROM {table_name} WHERE {row_name}=:entity_id'
    execute_query(db_session, delete_query, {'entity_id': entity_id})

    return None


def wait_for_data_in_db(
    db_session: Any,
    table_name: str,
    conditions: Dict[str, Any],
    timeout: int = 30,
    poll_interval: int = 5,
) -> Optional[Dict[str, Any]]:
    """Wait for the data to appear in the database with retries, starting with an initial wait."""
    # Initial wait before starting the polling
    time.sleep(poll_interval)

    start_time = time.time()

    while time.time() - start_time < timeout:
        db_result = get_entity(db_session, table_name, conditions)
        if db_result:
            return db_result
        time.sleep(poll_interval)  # Wait for a few seconds before retrying

    raise TimeoutError(
        f'Data not found in {table_name} table within {timeout} seconds.'
    )


def wait_for_generate_payments_to_complete(
    db_session: Any,
    payment_group_payment_id: int,
    timeout: int = 30,
    poll_interval: int = 5,
) -> Optional[Dict[str, Any]]:
    """Wait for generate_payments to complete, starting with an initial wait."""
    # Initial wait before starting the polling
    time.sleep(poll_interval)

    start_time = time.time()

    conditions = {
        'parent_table_name': 'payment_group_payment',
        'action_name': 'generate_payments',
        'parent_table_id': payment_group_payment_id,
    }

    while time.time() - start_time < timeout:
        db_result = get_entity(db_session, 'abacus_state', conditions)
        if db_result:
            action_status = db_result['action_status'].lower()
            log.info(f'Current status: {action_status}')
            if 'complete' == action_status:
                return db_result
        time.sleep(poll_interval)  # Wait for a few seconds before retrying

    raise TimeoutError(
        f'Record in abacus_state table with {conditions} did not reach complete within {timeout} seconds.'
    )