"""
SQLModel models generator.

Generates SQLModel table classes from parsed Django models.
"""

from typing import Any

from ..base import BaseGenerator, GeneratorContext
from ...ir.models import GeneratedFile, ParsedModel, ParsedField, RelationType
from ...parser.type_mapper import TypeMapper

# Reserved attribute names in SQLAlchemy/SQLModel that cannot be used as field names
RESERVED_FIELD_NAMES = {
    "metadata",  # SQLAlchemy MetaData object
    "registry",  # SQLAlchemy registry
    "query",     # SQLAlchemy query attribute
    "query_class",
}

# Type names that might conflict with field names - map to aliases
# When a field has the same name as its type, we need to alias the type import
TYPE_ALIASES = {
    "date": "DateType",
    "datetime": "DateTimeType",
    "time": "TimeType",
    "bytes": "BytesType",
    "type": "TypeType",
}


class SQLModelGenerator(BaseGenerator):
    """
    Generates SQLModel model classes.

    Creates one models.py per app with all table definitions.
    """

    def __init__(self, context: GeneratorContext):
        super().__init__(context)
        self.type_mapper = TypeMapper(
            use_jsonb=self.config.use_jsonb,
            use_array_fields=self.config.use_array_fields,
            use_uuid_type=self.config.use_uuid_type,
        )
        # Track if we need Column import (for sa_column usage)
        self._needs_column_import = False
        # Track if we need Index/UniqueConstraint imports
        self._needs_index_import = False
        self._needs_unique_constraint_import = False
        # Track if we need desc() for descending indexes
        self._needs_desc_import = False
        # Track if we need Optional import (for relationships)
        self._needs_optional_import = False
        # Track type aliases needed (when field name = type name)
        self._type_aliases_needed: set[str] = set()
        # Track enum classes used (from choices_class) — imported from enums.py
        self._enum_classes_used: set[str] = set()

    def generate(self) -> list[GeneratedFile]:
        """Generate SQLModel files for all apps."""
        files: list[GeneratedFile] = []

        for app_label, models in self.context.group_models_by_app().items():
            content = self._generate_models_file(app_label, models)
            files.append(self.create_file(f"{app_label}/models.py", content))

            # Create __init__.py for the app
            init_content = self._generate_init_file(models)
            files.append(self.create_file(f"{app_label}/__init__.py", init_content))

        return files

    def _generate_models_file(self, app_label: str, models: list[ParsedModel]) -> str:
        """Generate models.py content for an app."""
        self.type_mapper.reset_imports()
        self._needs_column_import = False
        self._needs_index_import = False
        self._needs_unique_constraint_import = False
        self._needs_desc_import = False
        self._needs_optional_import = False
        self._type_aliases_needed = set()

        # Reset enum tracking per app
        self._enum_classes_used = set()

        # Process all models first to collect imports
        model_definitions = []
        for model in models:
            model_def = self._generate_model_class(model)
            model_definitions.append(model_def)

        # Build imports section
        imports = self._build_imports()

        # Combine all parts
        parts = [
            '"""',
            f"SQLModel models for {app_label}.",
            "",
            "Auto-generated by django-cfg FastAPI ORM Generator.",
            "Do not edit manually - changes will be overwritten.",
            '"""',
            "",
            imports,
            "",
        ]

        parts.extend(model_definitions)

        return "\n".join(parts)

    def _generate_model_class(self, model: ParsedModel) -> str:
        """Generate a single SQLModel class."""
        lines = []

        # Class definition
        lines.append(f"class {model.name}(SQLModel, table=True):")

        # Docstring
        if self.config.add_docstrings:
            doc = model.verbose_name or f"{model.name} model"
            lines.append(f'    """{doc}."""')
            lines.append("")

        # Table name
        lines.append(f'    __tablename__ = "{model.table_name}"')
        lines.append("")

        # Generate fields
        for field in model.fields:
            # Skip M2M fields (they become relationships, not columns)
            if field.relation_type == RelationType.MANY_TO_MANY:
                continue

            field_line = self._generate_field(field)
            if field_line:
                lines.append(f"    {field_line}")

        # Add relationships for FK/O2O (optional - can cause issues with duplicate model names)
        if self.config.include_relationships and model.relationships:
            lines.append("")
            lines.append("    # Relationships")
            for rel in model.relationships:
                rel_line = self._generate_relationship(rel)
                if rel_line:
                    lines.append(f"    {rel_line}")

        # Add table args for constraints
        table_args = self._generate_table_args(model)
        if table_args:
            lines.append("")
            lines.append(f"    __table_args__ = {table_args}")

        lines.append("")
        lines.append("")

        return "\n".join(lines)

    def _generate_field(self, field: ParsedField) -> str:
        """Generate a single field definition."""
        # Get Python type
        python_type = self.type_mapper.get_python_type(field)

        # If field has a TextChoices enum class, use it as the type
        if field.choices_class:
            self._enum_classes_used.add(field.choices_class)
            if field.nullable:
                python_type = f"Optional[{field.choices_class}]"
                self.type_mapper._imports.add(("typing", "Optional"))
            else:
                python_type = field.choices_class

        # Handle field name conflicting with type name (e.g., date: date)
        # Need to use aliased type (e.g., date: DateType)
        if field.name in TYPE_ALIASES:
            # Check if the type annotation contains the conflicting type
            type_name = field.name
            if type_name in python_type:
                alias = TYPE_ALIASES[type_name]
                python_type = python_type.replace(type_name, alias)
                self._type_aliases_needed.add(type_name)

        # Handle reserved field names (like 'metadata' which conflicts with SQLAlchemy)
        attr_name = field.name
        db_column_name = None
        if field.name in RESERVED_FIELD_NAMES:
            attr_name = f"{field.name}_"
            db_column_name = field.name  # Need to map to original column name

        # Handle blank=True for string fields (make Optional with empty default)
        is_string_field = field.django_type in (
            'CharField', 'TextField', 'SlugField', 'EmailField', 'URLField', 'FilePathField'
        )
        if field.blank and is_string_field and not field.nullable:
            # blank=True without null=True means empty string is allowed
            if not python_type.startswith("Optional"):
                python_type = f"Optional[{python_type}]"
                self.type_mapper._imports.add(("typing", "Optional"))

        # Check if we need sa_column first (affects primary_key handling)
        sa_column = self.type_mapper.get_sqlalchemy_column(field)

        # If field uses a StrEnum (choices_class), force sa_column=Column(String(...)) to prevent
        # SQLAlchemy from creating a PostgreSQL ENUM type. Django stores these as plain VARCHAR.
        if field.choices_class and not sa_column:
            str_type = f"String({field.max_length})" if field.max_length else "String"
            sa_column = f"Column({str_type})"
            self.type_mapper._imports.add(("sqlalchemy", "String"))
            self._needs_column_import = True

        # If field name is reserved, we must use sa_column to specify the actual column name
        if db_column_name and not sa_column:
            # Create a basic sa_column with just the column name
            sa_column = f'Column("{db_column_name}")'

        # Get Field() definition parts
        field_parts = []

        # Default value and primary key
        # Note: SQLModel doesn't allow primary_key=True in Field() when sa_column is used
        # In that case, we need to put primary_key=True inside Column()
        if field.primary_key:
            field_parts.append("default=None")
            if not sa_column:
                # Only add primary_key to Field() if no sa_column
                field_parts.append("primary_key=True")
        elif field.default is not None:
            default_val = self.type_mapper.get_default_value(field)
            if default_val:
                if default_val.startswith("default_factory"):
                    field_parts.append(default_val)
                else:
                    field_parts.append(f"default={default_val}")
        elif field.nullable:
            field_parts.append("default=None")
        elif field.blank and is_string_field:
            # blank=True without null → default to empty string
            field_parts.append('default=""')

        # Foreign key
        if field.relation_type == RelationType.FOREIGN_KEY:
            table_name = field.related_table_name or (field.related_model_name.lower() if field.related_model_name else "unknown")
            field_parts.append(f'foreign_key="{table_name}.id"')

        # Constraints
        # Note: When sa_column is used, SQLModel requires constraints (unique, index, nullable)
        # to be inside Column(), not in Field()
        if not sa_column:
            if field.max_length and not field.is_relation:
                field_parts.append(f"max_length={field.max_length}")
            if field.unique and not field.primary_key:
                field_parts.append("unique=True")
            if field.db_index and not field.primary_key and not field.unique:
                field_parts.append("index=True")

        # Description
        if self.config.add_docstrings and field.help_text:
            help_text = field.help_text.replace('"', "'")[:100]
            field_parts.append(f'description="{help_text}"')

        # Add sa_column with constraints moved inside Column()
        if sa_column:
            # Collect Column() options
            column_opts = []
            if field.primary_key:
                column_opts.append("primary_key=True")
            if field.unique and not field.primary_key:
                column_opts.append("unique=True")
            if field.db_index and not field.primary_key and not field.unique:
                column_opts.append("index=True")
            if field.nullable:
                column_opts.append("nullable=True")

            # For reserved names, we need to add the column name as first argument
            if db_column_name:
                # Insert column name at the beginning: Column(TYPE) -> Column("name", TYPE)
                # or Column("name") -> Column("name", ...) with options
                if sa_column.startswith('Column("'):
                    # Already has column name (from basic sa_column created above)
                    if column_opts:
                        last_paren = sa_column.rfind(")")
                        sa_column = sa_column[:last_paren] + ", " + ", ".join(column_opts) + ")"
                else:
                    # Need to insert column name: Column(TYPE) -> Column("name", TYPE)
                    sa_column = sa_column.replace("Column(", f'Column("{db_column_name}", ', 1)
                    if column_opts:
                        last_paren = sa_column.rfind(")")
                        sa_column = sa_column[:last_paren] + ", " + ", ".join(column_opts) + ")"
            else:
                # Insert options into Column() before final )
                if column_opts:
                    last_paren = sa_column.rfind(")")
                    sa_column = sa_column[:last_paren] + ", " + ", ".join(column_opts) + ")"

            field_parts.append(f"sa_column={sa_column}")
            self._needs_column_import = True

        # Build Field() call
        if field_parts:
            field_def = f"Field({', '.join(field_parts)})"
        else:
            field_def = "Field()"

        return f"{attr_name}: {python_type} = {field_def}"

    def _generate_relationship(self, rel) -> str:
        """Generate a Relationship() definition."""
        type_hint = rel.get_type_hint()
        rel_def = rel.get_relationship_def()
        # Track if Optional is used in type hint
        if "Optional" in type_hint:
            self._needs_optional_import = True
        return f"{rel.name}: {type_hint} = {rel_def}"

    def _generate_table_args(self, model: ParsedModel) -> str:
        """Generate __table_args__ for constraints.

        Always includes extend_existing=True to allow table redefinition
        (needed for pytest test collection and hot reloading).
        """
        args = []

        # Unique together constraints
        for fields in model.unique_together:
            field_names = ", ".join(f'"{f}"' for f in fields)
            args.append(f"UniqueConstraint({field_names})")
            self._needs_unique_constraint_import = True

        # Indexes
        for idx in model.indexes:
            if idx.get('fields'):
                # Handle Django's -field syntax for DESC ordering
                field_parts = []
                for f in idx['fields']:
                    if f.startswith('-'):
                        # DESC ordering: -created_at -> desc("created_at")
                        col_name = f[1:]  # Strip the - prefix
                        field_parts.append(f'desc("{col_name}")')
                        self._needs_desc_import = True
                    else:
                        field_parts.append(f'"{f}"')

                field_names = ", ".join(field_parts)
                name = idx.get('name', '')
                if idx.get('unique'):
                    args.append(f'UniqueConstraint({field_names}, name="{name}")')
                    self._needs_unique_constraint_import = True
                else:
                    args.append(f'Index("{name}", {field_names})')
                    self._needs_index_import = True

        # Always add extend_existing to allow table redefinition in tests
        table_opts = '{"extend_existing": True}'

        if args:
            if len(args) == 1:
                return f"({args[0]}, {table_opts})"
            return f"(\n        {',\n        '.join(args)},\n        {table_opts},\n    )"

        return f"({table_opts},)"

    def _build_imports(self) -> str:
        """Build imports section based on collected imports."""
        lines = []

        # Standard library imports
        std_imports = []
        typing_imports = []
        datetime_imports = []

        for module, name in sorted(self.type_mapper.imports):
            if module == "typing":
                typing_imports.append(name)
            elif module == "datetime":
                datetime_imports.append(name)
            elif module in ("decimal", "uuid"):
                std_imports.append(f"from {module} import {name}")

        # Add Optional if needed for relationships
        if self._needs_optional_import:
            typing_imports.append("Optional")

        # Build datetime imports with aliases for conflicting types
        if datetime_imports:
            dt_parts = []
            for dt_name in sorted(set(datetime_imports)):
                if dt_name in self._type_aliases_needed:
                    dt_parts.append(f"{dt_name} as {TYPE_ALIASES[dt_name]}")
                else:
                    dt_parts.append(dt_name)
            lines.append(f"from datetime import {', '.join(dt_parts)}")

        if std_imports:
            lines.extend(sorted(set(std_imports)))
        if typing_imports:
            lines.append(f"from typing import {', '.join(sorted(set(typing_imports)))}")

        lines.append("")

        # SQLModel imports
        lines.append("from sqlmodel import SQLModel, Field, Relationship")

        # SQLAlchemy imports
        sa_imports = set()
        sa_pg_imports = []

        # Add Column if needed (for sa_column usage)
        if self._needs_column_import:
            sa_imports.add("Column")

        # Add Index/UniqueConstraint if needed
        if self._needs_index_import:
            sa_imports.add("Index")
        if self._needs_unique_constraint_import:
            sa_imports.add("UniqueConstraint")
        # Add desc() for descending indexes
        if self._needs_desc_import:
            sa_imports.add("desc")

        pgvector_imports = []

        for module, name in self.type_mapper.imports:
            if module == "sqlalchemy":
                sa_imports.add(name)
            elif module == "sqlalchemy.dialects.postgresql":
                sa_pg_imports.append(name)
            elif module == "pgvector.sqlalchemy":
                pgvector_imports.append(name)

        if sa_imports:
            lines.append(f"from sqlalchemy import {', '.join(sorted(sa_imports))}")
        if sa_pg_imports:
            lines.append(f"from sqlalchemy.dialects.postgresql import {', '.join(sorted(set(sa_pg_imports)))}")
        if pgvector_imports:
            lines.append(f"from pgvector.sqlalchemy import {', '.join(sorted(set(pgvector_imports)))}")

        # Enum imports from enums.py (same app)
        if self._enum_classes_used:
            lines.append(f"from .enums import {', '.join(sorted(self._enum_classes_used))}")

        return "\n".join(lines)

    def _generate_init_file(self, models: list[ParsedModel]) -> str:
        """Generate __init__.py with exports."""
        model_names = [m.name for m in models]
        imports = "\n".join(f"from .models import {name}" for name in model_names)

        # Check if any field in any model uses a TextChoices enum
        has_enums = any(
            f.choices_class
            for m in models
            for f in m.fields
            if f.choices_class
        )
        enum_import = "\nfrom .enums import *  # noqa: F401, F403" if has_enums else ""

        return f'''"""Models for this app."""

{imports}{enum_import}

__all__ = [{", ".join(f'"{n}"' for n in model_names)}]
'''
