from typing import Any, Type, Optional, List, Union, get_origin, get_args, ForwardRef
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from sqlalchemy import Integer, String, Boolean, TypeDecorator, DateTime, Date, ForeignKey as SAForeignKey, ForeignKeyConstraint as SAForeignKeyConstraint, UniqueConstraint as SAUniqueConstraint, CheckConstraint as SACheckConstraint, Index as SAIndex, Table, Column
import datetime
import uuid
from functools import wraps

_UNSET = object()

class Base(DeclarativeBase):
    pass

class Id:
    """Marker for primary key (@Id)"""
    def __init__(self, size: Optional[int] = None):
        self.size = size

class Size:
    """Marker for the maximum column size (Equivalent to @Size)"""
    def __init__(self, size: int = 256, nullable: Optional[bool] = None, default: Any = _UNSET):
        self.size = size
        self.nullable = nullable
        self.default = default

class Decimal:
    """Marker for decimal precision and scale (Equivalent to @Decimal)"""
    def __init__(self, precision: int, scale: int, nullable: Optional[bool] = None, default: Any = _UNSET):
        self.precision = precision
        self.scale = scale
        self.nullable = nullable
        self.default = default

class DateFormat(TypeDecorator):
    """
    Custom SQLAlchemy Type to handle Date formatting automatically.
    Converts strictly between Python strings/datetimes and Database DateTimes based on a format string.
    """
    impl = DateTime
    # Keeps it as a DateTime column when cached or processed by SQLite/PostgreSQL
    cache_ok = True

    def __init__(self, format: str = "%Y-%m-%d %H:%M:%S", *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.format = format

    def process_bind_param(self, value, dialect):
        # Convert String -> Datetime BEFORE saving to DB
        if value is None:
            return None
        if isinstance(value, str):
            try:
                # If we passed a string like "27/02/2026", Python converts it to Datetime object for SQLAlchemy
                return datetime.datetime.strptime(value, self.format)
            except ValueError:
                raise ValueError(f"Date string '{value}' does not match format '{self.format}'")
        return value # Returns native datetime if it's already one

    def process_result_value(self, value, dialect):
        # Convert Datetime -> String AFTER loading from DB
        # The user's Python class instance will see a formatted String if they desire, instead of a raw Datetime.
        if value is not None and isinstance(value, datetime.datetime):
            return value.strftime(self.format)
        return value

class DateFormatMarker:
    """Marker for date format properties (Equivalent to @JsonFormat or @DateTimeFormat)"""
    def __init__(self, format: str = "%Y-%m-%d %H:%M:%S", nullable: Optional[bool] = None, default: Any = _UNSET):
        self.format = format
        self.nullable = nullable
        self.default = default


class Unique:
    """Marker for unique constraints declared in @table(..., constraints=[...])."""
    def __init__(self, *fields: str, name: Optional[str] = None):
        if not fields:
            raise ValueError("Unique requires at least one field.")
        self.fields = fields
        self.name = name


class ForeignKey:
    """Marker for FK constraints declared in @table(..., constraints=[...])."""
    def __init__(
        self,
        local_fields: Union[str, List[str]],
        remote_fields: Union[str, List[str]],
        name: Optional[str] = None,
        ondelete: Optional[str] = None,
        onupdate: Optional[str] = None,
    ):
        self.local_fields = local_fields
        self.remote_fields = remote_fields
        self.name = name
        self.ondelete = ondelete
        self.onupdate = onupdate


class Check:
    """Marker for CHECK constraints declared in @table(..., constraints=[...])."""
    def __init__(self, expression: str, name: Optional[str] = None):
        if not expression or not str(expression).strip():
            raise ValueError("Check requires a non-empty expression.")
        self.expression = expression
        self.name = name


class Index:
    """Marker for index declarations in @table(..., constraints=[...])."""
    def __init__(self, *fields: str, name: Optional[str] = None, unique: bool = False):
        if not fields:
            raise ValueError("Index requires at least one field.")
        self.fields = fields
        self.name = name
        self.unique = unique

class ManyToOne:
    """Marker for many-to-one relationships with support for simple/composite foreign keys."""
    def __init__(self, fields: Union[str, List[str]], nullable: bool = True, back_populates: Optional[str] = None):
        self.fields = fields
        self.nullable = nullable
        self.back_populates = back_populates

class OneToOne:
    """Marker for one-to-one relationships with support for simple/composite foreign keys."""
    def __init__(self, fields: Union[str, List[str]], nullable: bool = True, back_populates: Optional[str] = None):
        self.fields = fields
        self.nullable = nullable
        self.back_populates = back_populates

class OneToMany:
    """Marker for one-to-many relationships. Use list annotation in the attribute type."""
    def __init__(self, mapped_by: Optional[str] = None):
        self.mapped_by = mapped_by

class ManyToMany:
    """Marker for many-to-many relationships. Use list annotation in the attribute type."""
    def __init__(self, mapped_by: Optional[str] = None, table_name: Optional[str] = None):
        self.mapped_by = mapped_by
        self.table_name = table_name


_association_tables: dict[str, Table] = {}


def _normalize_fields(fields: Union[str, List[str]]) -> list[str]:
    if isinstance(fields, str):
        parsed = [item.strip() for item in fields.split(",") if item.strip()]
    else:
        parsed = [str(item).strip() for item in fields if str(item).strip()]
    if not parsed:
        raise ValueError("Relationship fields cannot be empty.")
    return parsed


def _resolve_target_from_annotation(attr_type: Any):
    origin = get_origin(attr_type)
    if origin in (list,):
        args = get_args(attr_type)
        return args[0] if args else None
    return attr_type


def _extract_base_type_and_nullable(type_hint: Any) -> tuple[Any, bool]:
    origin = get_origin(type_hint)
    if origin is Union:
        args = get_args(type_hint)
        non_none_args = [arg for arg in args if arg is not type(None)]
        if len(non_none_args) == 1 and len(non_none_args) != len(args):
            return non_none_args[0], True
    return type_hint, False


def _resolve_target_name(target: Any) -> Optional[str]:
    if isinstance(target, ForwardRef):
        return target.__forward_arg__
    if isinstance(target, str):
        return target
    return None


def _build_secondary_table(left_cls: Type, right_cls: Type, custom_name: Optional[str] = None) -> Table:
    if not hasattr(left_cls, "__table__") or not hasattr(right_cls, "__table__"):
        raise ValueError("ManyToMany requires both sides to be mapped classes.")

    left_table = left_cls.__table__
    right_table = right_cls.__table__

    table_name = custom_name
    if not table_name:
        names = sorted([left_table.name, right_table.name])
        table_name = f"{names[0]}_{names[1]}"

    if table_name in _association_tables:
        return _association_tables[table_name]

    columns = []
    for pk_col in left_table.primary_key.columns:
        columns.append(
            Column(
                f"{left_table.name}_{pk_col.name}",
                pk_col.type,
                SAForeignKey(f"{left_table.name}.{pk_col.name}"),
                primary_key=True,
            )
        )

    for pk_col in right_table.primary_key.columns:
        columns.append(
            Column(
                f"{right_table.name}_{pk_col.name}",
                pk_col.type,
                SAForeignKey(f"{right_table.name}.{pk_col.name}"),
                primary_key=True,
            )
        )

    association = Table(table_name, Base.metadata, *columns)
    _association_tables[table_name] = association
    return association


def _coerce_table_constraints(constraints: Optional[List[Any]]) -> list[Any]:
    if not constraints:
        return []

    coerced = []
    for constraint in constraints:
        if isinstance(constraint, Unique):
            coerced.append(SAUniqueConstraint(*constraint.fields, name=constraint.name))
        elif isinstance(constraint, ForeignKey):
            local_cols = _normalize_fields(constraint.local_fields)
            remote_cols = _normalize_fields(constraint.remote_fields)
            if len(local_cols) != len(remote_cols):
                raise ValueError("ForeignKey local_fields and remote_fields must have the same number of columns.")
            coerced.append(
                SAForeignKeyConstraint(
                    local_cols,
                    remote_cols,
                    name=constraint.name,
                    ondelete=constraint.ondelete,
                    onupdate=constraint.onupdate,
                )
            )
        elif isinstance(constraint, Check):
            coerced.append(SACheckConstraint(constraint.expression, name=constraint.name))
        elif isinstance(constraint, Index):
            index_name = constraint.name or f"ix_{'_'.join(constraint.fields)}"
            coerced.append(SAIndex(index_name, *constraint.fields, unique=constraint.unique))
        else:
            coerced.append(constraint)
    return coerced


def _merge_table_args(existing_table_args: Any, generated_constraints: list[Any]) -> Any:
    if not generated_constraints:
        return existing_table_args

    if existing_table_args is None:
        return tuple(generated_constraints)

    if isinstance(existing_table_args, dict):
        return tuple(generated_constraints) + (existing_table_args,)

    if isinstance(existing_table_args, tuple):
        if existing_table_args and isinstance(existing_table_args[-1], dict):
            return tuple(existing_table_args[:-1]) + tuple(generated_constraints) + (existing_table_args[-1],)
        return tuple(existing_table_args) + tuple(generated_constraints)

    return (existing_table_args,) + tuple(generated_constraints)

def table(name: str, constraints: Optional[List[Any]] = None):
    """
    It intercepts the annotated class and returns a compatible SQLAlchemy class.
    """
    def decorator(cls: Type) -> Type[Base]:
        annotations = getattr(cls, '__annotations__', {})
        user_constraints = _coerce_table_constraints(constraints)
        attrs = {
            '__tablename__': name,
            '__module__': cls.__module__,
            '__doc__': cls.__doc__
        }
        table_constraints = []
        relationship_specs = []
        
        # Preserve user-defined methods and non-mapped class attributes
        for key, value in cls.__dict__.items():
            if key in attrs:
                continue
            if key in annotations:
                continue
            if key in ('__dict__', '__weakref__', '__slots__', '__annotations__'):
                continue
            attrs[key] = value
        
        for attr_name, attr_type in annotations.items():
            # Checks the value assigned to the property
            attr_val = getattr(cls, attr_name, None)
            has_explicit_default = attr_name in cls.__dict__
            resolved_attr_type, is_optional_type = _extract_base_type_and_nullable(attr_type)

            if isinstance(attr_val, (ManyToOne, OneToOne, OneToMany, ManyToMany)):
                target = _resolve_target_from_annotation(resolved_attr_type)

                if isinstance(attr_val, (ManyToOne, OneToOne)):
                    target_name = _resolve_target_name(target)
                    if target_name:
                        raise ValueError(
                            f"{attr_name}: ManyToOne/OneToOne requires a class type annotation, not forward reference."
                        )
                    if target is None:
                        raise ValueError(f"{attr_name}: relationship target could not be resolved.")

                    target_table = getattr(target, "__table__", None)
                    if target_table is None:
                        raise ValueError(f"{attr_name}: target entity must be decorated with @table before use.")

                    fields = _normalize_fields(attr_val.fields)
                    local_cols = []
                    remote_cols = []
                    for field_name in fields:
                        if field_name not in target_table.c:
                            raise ValueError(
                                f"{attr_name}: target field '{field_name}' does not exist in '{target.__name__}'."
                            )
                        remote_cols.append(field_name)
                        local_col_name = f"{attr_name}_{field_name}"
                        local_cols.append(local_col_name)

                        remote_col = target_table.c[field_name]
                        attrs[local_col_name] = mapped_column(remote_col.type, nullable=attr_val.nullable)

                    table_constraints.append(
                        SAForeignKeyConstraint(
                            local_cols,
                            [f"{target_table.name}.{col_name}" for col_name in remote_cols],
                        )
                    )

                    if isinstance(attr_val, OneToOne):
                        table_constraints.append(SAUniqueConstraint(*local_cols))

                    relationship_specs.append(
                        {
                            "kind": "one_to_one" if isinstance(attr_val, OneToOne) else "many_to_one",
                            "attr_name": attr_name,
                            "target": target,
                            "local_cols": local_cols,
                            "back_populates": attr_val.back_populates,
                        }
                    )

                elif isinstance(attr_val, OneToMany):
                    # Prefer target name so SQLAlchemy can resolve lazily from registry.
                    # This avoids failures when a concrete class object exists but is not yet mapped
                    # due to import timing/circular references.
                    relationship_specs.append(
                        {
                            "kind": "one_to_many",
                            "attr_name": attr_name,
                            "target": target,
                            "target_name": _resolve_target_name(target) or getattr(target, "__name__", None),
                            "mapped_by": attr_val.mapped_by,
                        }
                    )

                elif isinstance(attr_val, ManyToMany):
                    relationship_specs.append(
                        {
                            "kind": "many_to_many",
                            "attr_name": attr_name,
                            "target": target,
                            "target_name": _resolve_target_name(target),
                            "mapped_by": attr_val.mapped_by,
                            "table_name": attr_val.table_name,
                        }
                    )

                continue
            
            # Map Python native types to SQLAlchemy types
            sa_type = None
            if resolved_attr_type == int:
                sa_type = Integer
            elif resolved_attr_type == bool:
                sa_type = Boolean
            elif resolved_attr_type == str:
                sa_type = String(256)
            elif resolved_attr_type == float:
                from sqlalchemy import Float
                sa_type = Float
            elif resolved_attr_type == datetime.datetime:
                sa_type = DateTime
            elif resolved_attr_type == datetime.date:
                sa_type = Date

            # Is it a primary key?
            if isinstance(attr_val, Id):
                if resolved_attr_type == str:
                    if attr_val.size is not None:
                        # If it is a String and has a explicit size set in the Id
                        sa_type = String(attr_val.size)
                        attrs[attr_name] = mapped_column(sa_type, primary_key=True)
                    else:
                        attrs[attr_name] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
                else:
                    attrs[attr_name] = mapped_column(sa_type, primary_key=True)
            
            # Does it have a specific size?
            elif isinstance(attr_val, Size):
                # When using Size, we can also configure if it's part of a composite primary key
                is_pk = getattr(attr_val, 'primary_key', False)
                column_kwargs = {"primary_key": is_pk}
                if attr_val.nullable is not None:
                    column_kwargs["nullable"] = attr_val.nullable
                if attr_val.default is not _UNSET:
                    column_kwargs["default"] = attr_val.default
                attrs[attr_name] = mapped_column(String(attr_val.size), **column_kwargs)
                
            # Is it a decimal with precision?
            elif isinstance(attr_val, Decimal):
                from sqlalchemy import Numeric
                is_pk = getattr(attr_val, 'primary_key', False)
                column_kwargs = {"primary_key": is_pk}
                if attr_val.nullable is not None:
                    column_kwargs["nullable"] = attr_val.nullable
                if attr_val.default is not _UNSET:
                    column_kwargs["default"] = attr_val.default
                attrs[attr_name] = mapped_column(Numeric(attr_val.precision, attr_val.scale), **column_kwargs)
                
            # Is it formatted as Date? (Solves Both: string mapped to datetime or datetime mapped to string)
            elif isinstance(attr_val, DateFormatMarker):
                is_pk = getattr(attr_val, 'primary_key', False)
                column_kwargs = {"primary_key": is_pk}
                if attr_val.nullable is not None:
                    column_kwargs["nullable"] = attr_val.nullable
                if attr_val.default is not _UNSET:
                    column_kwargs["default"] = attr_val.default
                attrs[attr_name] = mapped_column(DateFormat(format=attr_val.format), **column_kwargs)
                
            # Normal attribute (e.g., age: int) without a custom initializer
            else:
                if sa_type:
                    column_kwargs = {}
                    if is_optional_type:
                        column_kwargs["nullable"] = True
                    if has_explicit_default:
                        column_kwargs["default"] = attr_val
                    attrs[attr_name] = mapped_column(sa_type, **column_kwargs)
                else:
                    column_kwargs = {}
                    if is_optional_type:
                        column_kwargs["nullable"] = True
                    if has_explicit_default:
                        column_kwargs["default"] = attr_val
                    attrs[attr_name] = mapped_column(**column_kwargs)

        combined_constraints = table_constraints + user_constraints
        attrs["__table_args__"] = _merge_table_args(attrs.get("__table_args__"), combined_constraints)
        
        # Creates a new class derived from SQLAlchemy DeclarativeBase
        sqla_class = type(cls.__name__, (Base,), attrs)

        for spec in relationship_specs:
            kind = spec["kind"]

            if kind in ("many_to_one", "one_to_one"):
                relation_kwargs = {
                    "foreign_keys": [getattr(sqla_class, col_name) for col_name in spec["local_cols"]]
                }
                if spec["back_populates"]:
                    relation_kwargs["back_populates"] = spec["back_populates"]
                if kind == "one_to_one":
                    relation_kwargs["uselist"] = False
                setattr(sqla_class, spec["attr_name"], relationship(spec["target"], **relation_kwargs))

            elif kind == "one_to_many":
                target = spec["target_name"] or spec["target"]
                relation_kwargs = {}
                if spec["mapped_by"]:
                    relation_kwargs["back_populates"] = spec["mapped_by"]
                setattr(sqla_class, spec["attr_name"], relationship(target, **relation_kwargs))

            elif kind == "many_to_many":
                target = spec["target"]
                if target is None or isinstance(target, (ForwardRef, str)) or not hasattr(target, "__table__"):
                    raise ValueError(
                        f"{spec['attr_name']}: ManyToMany requires a concrete target class annotation (e.g., list[Role])."
                    )
                secondary = _build_secondary_table(sqla_class, target, spec["table_name"])
                relation_kwargs = {"secondary": secondary}
                if spec["mapped_by"]:
                    relation_kwargs["back_populates"] = spec["mapped_by"]
                setattr(sqla_class, spec["attr_name"], relationship(target, **relation_kwargs))

        return sqla_class
        
    return decorator


# --- Global Database Configuration for Repositories ---

class DatabaseContext:
    engine = None

def configure_database(engine):
    """Sets the global engine for the sqla-lite framework."""
    DatabaseContext.engine = engine


def query(func):
    """Marks a repository method to run inside an auto-managed session query context."""
    setattr(func, "_sqla_lite_query", True)
    return func


def repository(entity_class: Type):
    """
    Class decorator equivalent to Spring's @Repository(Entity).
    Injects standard database operations methods transparently.
    """
    def decorator(cls: Type):
        def __init__(self, engine=None):
            self.engine = engine or DatabaseContext.engine
            if not self.engine:
                raise Exception("No database engine configured. Use 'configure_database(engine)' or pass it to the constructor.")
        
        def save(self, entity):
            from sqlalchemy.orm import Session
            with Session(self.engine, expire_on_commit=False) as session:
                session.add(entity)
                session.commit()
                return entity

        def get(self, *ident):
            from sqlalchemy.orm import Session
            with Session(self.engine, expire_on_commit=False) as session:
                # If multiple arguments are provided, pass them as a tuple (for composite keys)
                # If only one argument is provided, pass it directly (for simple keys)
                query_ident = ident[0] if len(ident) == 1 else ident
                return session.get(entity_class, query_ident)
                
        def delete(self, entity):
            from sqlalchemy.orm import Session
            with Session(self.engine, expire_on_commit=False) as session:
                entity = session.merge(entity)
                session.delete(entity)
                session.commit()
                
        def find_all(self):
            from sqlalchemy.orm import Session
            with Session(self.engine, expire_on_commit=False) as session:
                return session.query(entity_class).all()

        def _wrap_query_method(method):
            @wraps(method)
            def wrapped(self, *args, **kwargs):
                from sqlalchemy.orm import Session
                with Session(self.engine, expire_on_commit=False) as session:
                    query_ctx = session.query(self.entity_class)
                    return method(self, query_ctx, *args, **kwargs)
            return wrapped

        # Injects the base methods into the Repository class
        setattr(cls, '__init__', __init__)
        setattr(cls, 'save', save)
        setattr(cls, 'get', get)
        setattr(cls, 'delete', delete)
        setattr(cls, 'find_all', find_all)
        setattr(cls, 'entity_class', entity_class)

        for attr_name, attr_value in list(cls.__dict__.items()):
            if callable(attr_value) and getattr(attr_value, "_sqla_lite_query", False):
                setattr(cls, attr_name, _wrap_query_method(attr_value))
        
        return cls
    return decorator
