from typing import Literal, Callable, Any
from functools import wraps
import inspect
import time
import asyncio

from common.scheduling.utils import get_job_parameters
from common.scheduling.queue_utils import (
    attempt_lock_capture_with_token,
    release_lock_with_token,
)
from common.utils import format_string


def concurrency_lock(
    _func: Callable | None = None,
    *,
    # TTL
    lock_ttl_seconds: int = 600,
    # How long to wait trying to acquire the lock (0 = single attempt)
    acquire_timeout_seconds: float = 0.0,
    # New: retry interval between attempts while waiting
    acquire_retry_interval_seconds: float = 1.0,
    on_timeout: Literal["raise", "none"] = "raise",
    # Custom lock name template with variable substitution
    lock_name_template: str | None = None,
):
    """
    Ensures only one concurrent execution per lock key.

    Args:
        lock_ttl_seconds: Controls the lock's TTL (expiry in Redis).
        acquire_timeout_seconds: How long to try acquiring the lock (0 = single attempt).
        acquire_retry_interval_seconds: Retry interval between attempts while waiting.
        on_timeout: Action on timeout - "raise" to raise error, "none" to return None.
        lock_name_template: Optional template for custom lock names. Supports variables:
                           {module_name}, {class_name}, {method_name}, {instance_primary_key}.
                           Example: "my_lock:{class_name}.{method_name}:{instance_primary_key}"
                           Default: "method_lock:{module_name}.{class_name}.{method_name}:{instance_primary_key}"

    Notes:
      - Supports both @concurrency_lock and @concurrency_lock(...).
      - Template variables are automatically populated from method context.
      - {instance_primary_key} is only available when the method is bound to a SQLModel instance.
    """

    def decorator(method: Callable):
        if not inspect.iscoroutinefunction(method):
            raise TypeError("concurrency_lock requires an async function")

        @wraps(method)
        async def wrapper(*args: Any, **kwargs: Any):
            owner = args[0] if args else None

            # Prefer SQLModel metadata when available; otherwise fall back to owner class
            fn_parameters = get_job_parameters(
                method, ensure_sql_model=False, owner=owner
            )

            method_name = fn_parameters["method_name"]
            module_name = fn_parameters.get("module_name")
            model_name = fn_parameters.get("model_name")
            pk = fn_parameters.get("instance_primary_key", None)

            if module_name is None and owner is not None:
                module_name = getattr(owner.__class__, "__module__", "global")
            if model_name is None and owner is not None:
                model_name = getattr(owner.__class__, "__name__", "global")

            module_name = module_name or "global"
            model_name = model_name or "global"

            # Build lock name from template or default format
            template = (
                lock_name_template
                or "method_lock:{module_name}.{class_name}.{method_name}:{instance_primary_key}"
            )
            template_context = {
                "module_name": module_name,
                "class_name": model_name,
                "method_name": method_name,
                "instance_primary_key": pk or "None",
            }

            lock_name = format_string(template, **template_context)

            # Resolve TTL
            ttl_seconds = lock_ttl_seconds

            token: str | None = None

            if acquire_timeout_seconds <= 0:
                # Single-attempt mode
                captured, token = await attempt_lock_capture_with_token(
                    lock_name, expire_in_seconds=ttl_seconds
                )
                if not captured:
                    if on_timeout == "none":
                        return None
                    raise TimeoutError(f"Could not capture lock: {lock_name}")
            else:
                # Timed acquisition with retries
                deadline = time.monotonic() + acquire_timeout_seconds
                while True:
                    captured, token = await attempt_lock_capture_with_token(
                        lock_name, expire_in_seconds=ttl_seconds
                    )
                    if captured:
                        break
                    if time.monotonic() >= deadline:
                        if on_timeout == "none":
                            return None
                        raise TimeoutError(
                            f"Timed out acquiring lock after {acquire_timeout_seconds}s: {lock_name}"
                        )
                    await asyncio.sleep(acquire_retry_interval_seconds)

            try:
                return await method(*args, **kwargs)
            finally:
                # Only attempt to release if we actually captured and still own the lock
                if token is not None:
                    await release_lock_with_token(lock_name, token)

        return wrapper

    # Support bare decorator usage: @concurrency_lock
    if callable(_func):
        return decorator(_func)

    # Parameterized usage: @concurrency_lock(...)
    return decorator
