"""Shared helpers for sync and async CRUD operations."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, TypeVar

from pydynox._internal._atomic import AtomicOp, serialize_atomic
from pydynox._internal._conditions import ConditionPath
from pydynox.exceptions import ItemTooLargeException
from pydynox.hooks import HookType

if TYPE_CHECKING:
    from pydynox.client import DynamoDBClient
    from pydynox.conditions import Condition
    from pydynox.model import Model

M = TypeVar("M", bound="Model")

# Type aliases for return types
SaveResult = tuple[
    "DynamoDBClient",
    str,
    dict[str, Any],
    str | None,
    dict[str, str] | None,
    dict[str, Any] | None,
    bool,
]
SmartSaveResult = tuple[
    "DynamoDBClient",
    str,
    dict[str, Any],  # key or item
    str | None,  # condition_expr
    dict[str, str] | None,  # attr_names
    dict[str, Any] | None,  # attr_values
    bool,  # skip_hooks
    bool,  # use_update (True = UpdateItem, False = PutItem)
    dict[str, Any] | None,  # updates (for UpdateItem)
]
UpdateResult = tuple[
    "DynamoDBClient",
    str,
    dict[str, Any],
    str | None,
    str | None,
    dict[str, str] | None,
    dict[str, Any] | None,
    dict[str, Any] | None,
    bool,
]
UpdateByKeyResult = tuple[
    "DynamoDBClient",
    str,
    dict[str, Any],
    dict[str, Any],
    str | None,
    dict[str, str] | None,
    dict[str, Any] | None,
]
DeleteByKeyResult = tuple[
    "DynamoDBClient", str, dict[str, Any], str | None, dict[str, str] | None, dict[str, Any] | None
]


def prepare_get(
    cls: type[M], consistent_read: bool | None, keys: dict[str, Any]
) -> tuple[DynamoDBClient, str, dict[str, Any], bool]:
    """Prepare get operation. Returns (client, table, keys, use_consistent)."""
    client = cls._get_client()
    table = cls._get_table()

    # Translate python key names to DynamoDB aliases
    aliased_keys = {}
    for k, v in keys.items():
        dynamo_name = cls._py_to_dynamo.get(k, k)
        aliased_keys[dynamo_name] = v

    use_consistent = consistent_read
    if use_consistent is None:
        use_consistent = getattr(cls.model_config, "consistent_read", False)

    return client, table, aliased_keys, use_consistent


def finalize_get(cls: type[M], item: dict[str, Any] | None) -> M | None:
    """Finalize get operation - convert to model and run hooks."""
    if item is None:
        return None

    instance = cls.from_dict(item)
    skip = cls.model_config.skip_hooks if hasattr(cls, "model_config") else False
    if not skip:
        instance._run_hooks(HookType.AFTER_LOAD)
    return instance


def prepare_save(self: Model, condition: Condition | None, skip_hooks: bool | None) -> SaveResult:
    """Prepare save operation.

    Returns (client, table, item, condition_expr, attr_names, attr_values, skip).
    """
    skip = self._should_skip_hooks(skip_hooks)

    if not skip:
        self._run_hooks(HookType.BEFORE_SAVE)

    self._apply_auto_generate()

    version_attr = self._get_version_attr_name()
    version_condition, new_version = self._build_version_condition()

    final_condition = condition
    if version_condition is not None:
        final_condition = (
            final_condition & version_condition if final_condition else version_condition
        )

    if version_attr is not None:
        setattr(self, version_attr, new_version)

    max_size = (
        getattr(self.model_config, "max_size", None) if hasattr(self, "model_config") else None
    )
    if max_size is not None:
        size = self.calculate_size()
        if size.bytes > max_size:
            raise ItemTooLargeException(
                size=size.bytes,
                max_size=max_size,
                item_key=self._get_key(),
            )

    client = self._get_client()
    table = self._get_table()
    item = self.to_dict()

    if final_condition is not None:
        names: dict[str, str] = {}
        values: dict[str, Any] = {}
        expr = final_condition.serialize(names, values)
        attr_names = {v: k for k, v in names.items()}
        return client, table, item, expr, attr_names, values, skip

    return client, table, item, None, None, None, skip


def finalize_save(self: Model, skip: bool) -> None:
    """Finalize save operation - run after hooks and reset change tracking."""
    if not skip:
        self._run_hooks(HookType.AFTER_SAVE)
    # Reset change tracking after successful save
    self._reset_change_tracking()


def prepare_smart_save(
    self: Model, condition: Condition | None, skip_hooks: bool | None, full_replace: bool
) -> SmartSaveResult:
    """Prepare smart save - uses UpdateItem for changed fields only.

    Returns SmartSaveResult tuple with use_update flag indicating which operation to use.
    """
    skip = self._should_skip_hooks(skip_hooks)

    if not skip:
        self._run_hooks(HookType.BEFORE_SAVE)

    self._apply_auto_generate()

    version_attr = self._get_version_attr_name()
    version_condition, new_version = self._build_version_condition()

    final_condition = condition
    if version_condition is not None:
        final_condition = (
            final_condition & version_condition if final_condition else version_condition
        )

    if version_attr is not None:
        setattr(self, version_attr, new_version)

    max_size = (
        getattr(self.model_config, "max_size", None) if hasattr(self, "model_config") else None
    )
    if max_size is not None:
        size = self.calculate_size()
        if size.bytes > max_size:
            raise ItemTooLargeException(
                size=size.bytes,
                max_size=max_size,
                item_key=self._get_key(),
            )

    client = self._get_client()
    table = self._get_table()

    # Decide: UpdateItem (smart) or PutItem (full)
    # Use UpdateItem if: has original (loaded from DB), has changes, not full_replace
    use_update = self._original is not None and len(self._changed) > 0 and not full_replace

    if use_update:
        # Smart update: only send changed fields
        key = self._get_key()
        updates = {}
        for field_name in self._changed:
            if field_name in self._attributes:
                value = getattr(self, field_name, None)
                # Serialize the value - handles special types like S3Value
                attr = self._attributes[field_name]
                # Use alias for DynamoDB attribute name
                dynamo_name = self._py_to_dynamo.get(field_name, field_name)
                updates[dynamo_name] = attr.serialize(value)

        # Also include version field if it changed
        if version_attr is not None and version_attr not in updates:
            dynamo_version = self._py_to_dynamo.get(version_attr, version_attr)
            updates[dynamo_version] = new_version

        if final_condition is not None:
            names: dict[str, str] = {}
            values: dict[str, Any] = {}
            expr = final_condition.serialize(names, values)
            attr_names = {v: k for k, v in names.items()}
            # Rename condition value placeholders to avoid collision with update placeholders
            # Rust build_set_expression uses :v0, :v1, etc. for updates
            renamed_values: dict[str, Any] = {}
            renamed_expr = expr
            for old_key, val in values.items():
                new_key = old_key.replace(":v", ":cond")
                renamed_values[new_key] = val
                renamed_expr = renamed_expr.replace(old_key, new_key)
            return client, table, key, renamed_expr, attr_names, renamed_values, skip, True, updates

        return client, table, key, None, None, None, skip, True, updates

    # Full replace: PutItem with all fields
    item = self.to_dict()

    if final_condition is not None:
        names = {}
        values: dict[str, Any] = {}
        expr = final_condition.serialize(names, values)
        attr_names = {v: k for k, v in names.items()}
        return client, table, item, expr, attr_names, values, skip, False, None

    return client, table, item, None, None, None, skip, False, None


def prepare_delete(self: Model, condition: Condition | None, skip_hooks: bool | None) -> SaveResult:
    """Prepare delete. Returns (client, table, key, cond_expr, attr_names, attr_values, skip)."""
    skip = self._should_skip_hooks(skip_hooks)

    if not skip:
        self._run_hooks(HookType.BEFORE_DELETE)

    version_attr = self._get_version_attr_name()
    version_condition: Condition | None = None
    if version_attr is not None:
        current_version: int | None = getattr(self, version_attr, None)
        if current_version is not None:
            # Use alias for DynamoDB attribute name in condition
            dynamo_name = self._py_to_dynamo.get(version_attr, version_attr)
            path = ConditionPath(path=[dynamo_name])
            version_condition = path == current_version

    final_condition = condition
    if version_condition is not None:
        final_condition = (
            final_condition & version_condition if final_condition else version_condition
        )

    client = self._get_client()
    table = self._get_table()
    key = self._get_key()

    if final_condition is not None:
        names: dict[str, str] = {}
        values: dict[str, Any] = {}
        expr = final_condition.serialize(names, values)
        attr_names = {v: k for k, v in names.items()}
        return client, table, key, expr, attr_names, values, skip

    return client, table, key, None, None, None, skip


def finalize_delete(self: Model, skip: bool) -> None:
    """Finalize delete operation - run after hooks."""
    if not skip:
        self._run_hooks(HookType.AFTER_DELETE)


def prepare_update(
    self: Model,
    atomic: list[AtomicOp] | None,
    condition: Condition | None,
    skip_hooks: bool | None,
    kwargs: dict[str, Any],
) -> UpdateResult:
    """Prepare update. Returns UpdateResult tuple."""
    skip = self._should_skip_hooks(skip_hooks)

    if not skip:
        self._run_hooks(HookType.BEFORE_UPDATE)

    client = self._get_client()
    table = self._get_table()
    key = self._get_key()

    if atomic:
        update_expr, names, values = serialize_atomic(atomic)
        attr_names = {v: k for k, v in names.items()}

        cond_expr = None
        if condition is not None:
            cond_names: dict[str, str] = dict(names)
            cond_expr = condition.serialize(cond_names, values)
            cond_attr_names = {v: k for k, v in cond_names.items()}
            attr_names = {**attr_names, **cond_attr_names}

        return (
            client,
            table,
            key,
            update_expr,
            cond_expr,
            attr_names if attr_names else None,
            values if values else None,
            None,
            skip,
        )

    if kwargs:
        for attr_name, value in kwargs.items():
            if attr_name not in self._attributes:
                raise ValueError(f"Unknown attribute: {attr_name}")
            setattr(self, attr_name, value)

        # Translate update keys to DynamoDB alias names
        aliased_kwargs = {}
        for attr_name, value in kwargs.items():
            dynamo_name = self._py_to_dynamo.get(attr_name, attr_name)
            aliased_kwargs[dynamo_name] = value

        if condition is not None:
            cond_names: dict[str, str] = {}
            cond_values: dict[str, Any] = {}
            cond_expr = condition.serialize(cond_names, cond_values)
            attr_names = {v: k for k, v in cond_names.items()}
            return (
                client,
                table,
                key,
                None,
                cond_expr,
                attr_names,
                cond_values,
                aliased_kwargs,
                skip,
            )

        return client, table, key, None, None, None, None, aliased_kwargs, skip

    return client, table, key, None, None, None, None, None, skip


def finalize_update(self: Model, skip: bool) -> None:
    """Finalize update operation - run after hooks."""
    if not skip:
        self._run_hooks(HookType.AFTER_UPDATE)


def prepare_update_by_key(
    cls: type[M], condition: Condition | None, kwargs: dict[str, Any]
) -> UpdateByKeyResult | None:
    """Prepare update_by_key. Returns tuple or None if nothing to update."""
    key, updates = cls._extract_key_from_kwargs(kwargs)

    if not updates:
        return None

    for attr_name in updates:
        if attr_name not in cls._attributes:
            raise ValueError(f"Unknown attribute: {attr_name}")

    # Translate update keys to DynamoDB alias names
    aliased_updates = {}
    for attr_name, value in updates.items():
        dynamo_name = cls._py_to_dynamo.get(attr_name, attr_name)
        aliased_updates[dynamo_name] = value

    client = cls._get_client()
    table = cls._get_table()

    if condition is not None:
        names: dict[str, str] = {}
        values: dict[str, Any] = {}
        cond_expr = condition.serialize(names, values)
        attr_names = {v: k for k, v in names.items()}
        # Rename value placeholders to avoid collision
        renamed_values: dict[str, Any] = {}
        renamed_expr = cond_expr
        for old_key, val in values.items():
            new_key = old_key.replace(":v", ":cond")
            renamed_values[new_key] = val
            renamed_expr = renamed_expr.replace(old_key, new_key)
        return (
            client,
            table,
            key,
            aliased_updates,
            renamed_expr,
            attr_names if attr_names else None,
            renamed_values if renamed_values else None,
        )

    return client, table, key, aliased_updates, None, None, None


def prepare_delete_by_key(
    cls: type[M], condition: Condition | None, kwargs: dict[str, Any]
) -> DeleteByKeyResult:
    """Prepare delete_by_key. Returns (client, table, key, cond_expr, attr_names, attr_values)."""
    key, _ = cls._extract_key_from_kwargs(kwargs)

    client = cls._get_client()
    table = cls._get_table()

    if condition is not None:
        names: dict[str, str] = {}
        values: dict[str, Any] = {}
        cond_expr = condition.serialize(names, values)
        attr_names = {v: k for k, v in names.items()}
        return client, table, key, cond_expr, attr_names, values

    return client, table, key, None, None, None
