# decorators/subtask_handler.py

from functools import wraps
from typing import Callable, Any
from uuid import UUID
import json

from backend.app.db.functions.tasks.Tasks import SubTaskInterface


def subtask_handler():
    """
    Decorator that wraps your existing handler and automatically:

    - Extracts task_id + sub_task_id from event
    - Marks subtask as in progress
    - Executes handler logic
    - Marks subtask complete on success
    - Marks failed on exception
    """

    def decorator(func: Callable[..., Any]):

        @wraps(func)
        def wrapper(event: dict[str, Any], context: Any, *args, **kwargs):

            records = event.get("Records", [event])

            interface = SubTaskInterface()

            for record in records:

                # -------------------------------
                # Parse body safely
                # -------------------------------
                body = {}

                if isinstance(record.get("body"), str):
                    try:
                        body = json.loads(record["body"])
                    except Exception:
                        body = {}
                else:
                    body = record.get("body", {}) or {}

                task_id_raw = body.get("task_id")
                subtask_id_raw = body.get("sub_task_id")

                task_id = UUID(task_id_raw) if isinstance(task_id_raw, str) else None
                subtask_id = (
                    UUID(subtask_id_raw) if isinstance(subtask_id_raw, str) else None
                )

                if not task_id or not subtask_id:
                    raise RuntimeError("task_id or sub_task_id missing")

                # -------------------------------
                # Mark in progress
                # -------------------------------
                interface.update_subtask_status(
                    subtask_id=subtask_id,
                    status="in progress",
                )

                try:
                    # Pass the parsed body into your function
                    result = func(body, context, *args, **kwargs)

                    # -------------------------------
                    # Success → mark complete
                    # -------------------------------
                    interface.update_subtask_status(
                        subtask_id=subtask_id,
                        status="complete",
                        outputs={"result": result} if result else None,
                    )

                except Exception as e:

                    # -------------------------------
                    # Failure → mark failed
                    # -------------------------------
                    interface.update_subtask_status(
                        subtask_id=subtask_id,
                        status="failed",
                        outputs={"error": str(e)},
                    )

                    raise

            return None

        return wrapper

    return decorator