"""Custom sqlacodegen generator for Abacus models.

This generator customizes how SQLAlchemy models are generated from the database,
allowing for specific naming conventions, relationships, and model configurations
tailored to the Abacus team.
"""

import re
from collections.abc import Sequence
from typing import Any

from sqlacodegen.generators import DataclassGenerator
from sqlacodegen.models import ColumnAttribute, Model, ModelClass, RelationshipAttribute
from sqlalchemy import Column, Date, DateTime, MetaData, Table
from sqlalchemy.engine import Connection, Engine
from sqlalchemy.sql.expression import ColumnCollection

from .utils.strings import to_snake_case

MIXINS = {
    'AuditMixin': {'audit_id'},
    'CreateMixin': {'created_at', 'created_by'},
    'SoftDeleteMixin': {'deleted_at', 'deleted_by'},
    'UpdateMixin': {'last_modified', 'last_modified_by'},
}

MANAGED_COLUMNS = MIXINS['CreateMixin'] | MIXINS['UpdateMixin']


class AbacusGenerator(DataclassGenerator):
    """Custom generator for Abacus SQLAlchemy models."""

    def __init__(
        self,
        metadata: MetaData,
        bind: Connection | Engine,
        options: Sequence[str],
        *,
        base_module_path: str = 'abacus_models.core',
        **kwargs,
    ):
        """Initialize the generator.

        Args:
            metadata: SQLAlchemy MetaData object
            bind: Database connection or engine
            options: Generator options
            base_module_path: Module path for imports (default: 'abacus_models.core')
            **kwargs: Additional keyword arguments passed to parent class
        """
        super().__init__(metadata, bind, options, **kwargs)
        self.base_module_path: str = base_module_path
        self.mixins_module_path: str = f'{base_module_path}.mixins'
        self.types_module_path: str = f'{base_module_path}.types'

    def collect_imports_for_column(self, column: Column[Any]) -> None:
        """Determine each column's imports."""
        if isinstance(column.type, DateTime):
            # Replace DateTime with NormalizedDateTime
            # self.remove_literal_import('sqlalchemy', DateTime.__name__)
            self.add_literal_import(self.types_module_path, 'NormalizedDateTime')
        elif isinstance(column.type, Date):
            # Replace Date with NormalizedDate
            # self.remove_literal_import('sqlalchemy', Date.__name__)
            self.add_literal_import(self.types_module_path, 'NormalizedDate')
        else:
            # Use the default behavior
            super().collect_imports_for_column(column)

    def generate(self) -> str:
        """Generate the output."""
        # Get the database name
        db_name = 'Database'
        if isinstance(self.bind, Connection):
            db_name = self.bind.engine.url.database
        elif isinstance(self.bind, Engine):
            db_name = self.bind.url.database

        # Add a docstring
        generated = super().generate()
        generated = f'"""{db_name} Models."""\n\n' + generated

        # Return the output
        return generated

    def generate_base(self) -> None:
        """Generate the base class with mixins."""
        # Call the parent method to generate the base class
        super().generate_base()

        # Add BaseMixin import
        self.add_literal_import(self.mixins_module_path, 'BaseMixin')

        # Modify the base class declaration to include BaseMixin
        text = re.sub(r'\):$', ', BaseMixin):', self.base.declarations[0])
        self.base.declarations[0] = text

    def get_parent_class_names(self, model: Model) -> list[str]:
        """Determine parent classes for a model."""
        parents = ['Base']

        # Add mixins
        column_names: set[str] = {col.name.lower() for col in model.table.columns}
        for mixin_name, mixin_cols in MIXINS.items():
            if not column_names.isdisjoint(mixin_cols):
                self.add_literal_import(self.mixins_module_path, mixin_name)
                parents.append(mixin_name)

        return parents

    def get_sorted_column_attributes(
        self, attrs: list[ColumnAttribute]
    ) -> list[ColumnAttribute]:
        """Sort attributes to ensure non-default columns come before default columns."""
        def_attrs = []
        req_attrs = []

        for attr in attrs:
            column = attr.column
            if column.default or column.name in MANAGED_COLUMNS:
                def_attrs.append(attr)
            else:
                req_attrs.append(attr)

        return req_attrs + def_attrs

    def get_sorted_columns(self, columns: ColumnCollection) -> ColumnCollection:
        """Sort attributes to ensure non-default columns come before default columns."""
        collection = ColumnCollection()
        def_cols = []

        for column in columns:
            if column.default or column.name in MANAGED_COLUMNS:
                def_cols.append(column)
            else:
                collection.add(column)

        for column in def_cols:
            collection.add(column)

        return collection

    def render_class_declaration(self, model: Model) -> str:
        """Render the class declaration with custom parent classes."""
        parents = self.get_parent_class_names(model)
        parent_class_names = ', '.join(parents)
        return f'class {model.name}({parent_class_names}):'

    def render_models(self, models: list[Model]) -> str:
        """Render models."""
        rendered = [self.render_model(model) for model in models]
        return '\n\n\n'.join(rendered)

    def render_model(self, model: Model) -> str:
        """Render model with sorted columns."""
        columns = model.table.columns
        sorted_columns = self.get_sorted_columns(columns)

        if isinstance(model, ModelClass):
            attrs = model.columns
            sorted_attrs = self.get_sorted_column_attributes(attrs)
            try:
                model.columns = sorted_attrs
                model.table.c = sorted_columns  # Also updates .columns
                return self.render_class(model)
            finally:
                model.table.c = columns
                model.columns = attrs

        try:
            model.table.c = sorted_columns  # Also updates .columns
            return f'{model.name} = {self.render_table(model.table)}'
        finally:
            model.table.c = columns

    def render_column(
        self, column: Column[Any], show_name: bool, is_table: bool = False
    ) -> str:
        """Render column."""
        args: dict[str, Any] = {}

        # If the column is auto-increment
        ai_column = column.table.autoincrement_column
        if ai_column is not None and column.name == ai_column.name:
            # Make auto-increment
            args['autoincrement'] = True
            # Remove from __init__
            args['init'] = False

        # If the column is managed by a mixin
        if column.name in MANAGED_COLUMNS:
            # Make it optional in __init__
            args['default'] = None

        # Render the column
        rendered = super().render_column(column, show_name, is_table)

        # Ignore args that were already rendered
        args = {
            key: val
            for key, val in args.items()
            if not re.search(rf'[,(]\s*{key}\s*=', rendered)
        }

        # Convert args into a comma-separated 'key=value' string
        args_arr = [f', {key}={val}' for key, val in args.items()]
        args_str = ''.join(args_arr)

        # Append args to rendered column
        rendered = re.sub(r'\)$', f'{args_str})', rendered)

        return rendered

    def render_column_callable(self, is_table: bool, *args: Any, **kwargs: Any) -> str:
        """Render column callable."""
        is_nullable = kwargs.get('nullable', True)
        is_primary_key = 'primary_key' in kwargs

        # Make optional columns optional
        if not is_primary_key and is_nullable:
            kwargs.setdefault('default', None)

        return super().render_column_callable(is_table, *args, **kwargs)

    def render_column_type(self, column_type: Any) -> str:
        """Render the column type."""
        if isinstance(column_type, DateTime):
            return 'NormalizedDateTime'

        if isinstance(column_type, Date):
            return 'NormalizedDate'

        return super().render_column_type(column_type)

    def render_relationship(self, relationship: RelationshipAttribute) -> str:
        """Render relationship."""
        rendered = super().render_relationship(relationship)
        return re.sub(r'\)$', ', init=False)', rendered)

    def should_ignore_table(self, table: Table) -> bool:
        """Determine if a table should be ignored."""
        name = to_snake_case(table.name)

        # Skip rollbacks
        if 'rollback' in name:
            return True

        # Skip changelogs
        if 'databasechangelog' in name:
            return True

        # Skip temps
        if re.search(r'(?:^te?mp_)|(?:_te?mp$)', name):
            return True

        # Skip backfills
        if re.search(r'(?:^backfill_)|(?:_backfill$)', name):
            return True

        # Skip backups
        if re.search(r'(?:^backup_)|(?:_backup$)', name):
            return True

        # Skip tickets (e.g. ACC-9099)
        if re.search(r'(?:^|_)\w{2,8}_?\d{3,}(?:_|$)', name):
            return True

        # Skip dated (e.g. _20251011)
        if re.search(r'_\d{8}$', name):
            return True

        # Include
        return False