from typing import Any, Callable, List, Optional, Dict, Union, cast
from dataclasses import dataclass, field
import functools
import structlog
from datetime import timedelta

logger = structlog.get_logger()


@dataclass
class Retriever:
    name: str
    func: Callable[..., List[Dict[str, Any]]]
    backend: str = "custom"
    cache_ttl: Optional[timedelta] = None
    description: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __hash__(self) -> int:
        return hash(self.name)


class RetrieverRegistry:
    def __init__(self) -> None:
        self.retrievers: Dict[str, Retriever] = {}

    def register(self, retriever: Retriever) -> None:
        if retriever.name in self.retrievers:
            logger.warning(f"Overwriting existing retriever: {retriever.name}")
        self.retrievers[retriever.name] = retriever
        logger.info(f"Registered retriever: {retriever.name}")

    def get(self, name: str) -> Optional[Retriever]:
        return self.retrievers.get(name)


def retriever(
    backend: str = "custom",
    cache_ttl: Optional[Union[str, timedelta]] = None,
    name: Optional[str] = None,
) -> Any:
    """
    Decorator to register a function as a Retriever.

    Args:
        backend: "custom" (Python) or "postgres" (SQL).
        cache_ttl: Optional TTL for caching results.
        name: Optional override for retriever name.
    """

    def decorator(
        func: Callable[..., List[Dict[str, Any]]]
    ) -> Callable[..., List[Dict[str, Any]]]:
        # Resolve effective name
        r_name = name or func.__name__

        # Parse TTL
        parsed_ttl = None
        if cache_ttl:
            # Basic string parsing fallback if not a timedelta
            if isinstance(cache_ttl, timedelta):
                parsed_ttl = cache_ttl

        ret_obj = Retriever(
            name=r_name,
            func=func,
            backend=backend,
            cache_ttl=parsed_ttl,
            description=func.__doc__,
        )

        # Attach to function for inspection
        setattr(func, "_meridian_retriever", ret_obj)

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> List[Dict[str, Any]]:
            # Check for injected cache backend
            store_backend = getattr(ret_obj, "_cache_backend", None)

            # If caching is enabled and backend is available
            if ret_obj.cache_ttl and store_backend:
                try:
                    import json
                    import hashlib

                    # Simple cache key generation
                    # We hash args/kwargs.
                    # Note: this requires picklable/jsonable args.
                    key_parts = [r_name, str(args), str(kwargs)]
                    key_str = json.dumps(key_parts, sort_keys=True, default=str)
                    key_hash = hashlib.sha256(key_str.encode("utf-8")).hexdigest()
                    _cache_key = f"meridian:retriever:{r_name}:{key_hash}"

                    # Try get
                    # We assume store_backend is standard Redis client or similar interface
                    # Since wrapper is sync, but Redis is async... this is tricky.
                    # Retrievers might be async or sync.
                    # If func is async, wrapper should be async.
                    # For now, let's assume async if we want real async redis.
                    # For Phase 2, let's assume Retrievers are generic.
                    # If the user defined an async function, wrapper should be async.

                    # Handling Sync/Async wrapper transparency is complex.
                    # For Phase 2 MVP: Let's assume user calls `await func(...)` if it's async.
                    # If `func` is sync, we can't easily wait for async redis.

                    # DECISION: Retrievers should be ASYNC for optimal IO.
                    # "Story 2.1.2": Cache expensive vector search.
                    # Vector search is usually IO bound.
                    # Use `import inspect` to detect async.
                except Exception as e:
                    logger.warning(f"Cache key generation failed: {e}")

            from typing import cast

            return cast(List[Dict[str, Any]], list(func(*args, **kwargs)))

        # Helper for DAG Resolution
        async def _resolve_args(args: Any, kwargs: Any) -> Any:
            store_ref = getattr(ret_obj, "_meridian_store_ref", None)
            if not store_ref:
                return args, kwargs

            # Need entity_id to resolve
            # Convention: entity_id in kwargs or first arg?
            # Let's check kwargs first
            entity_id = kwargs.get("entity_id")
            if not entity_id:
                # Weak heuristic: check if first arg looks like entity_id? No, too risky.
                # If no entity_id, we can't resolve features.
                return args, kwargs

            from meridian.graph import DependencyResolver

            resolver = DependencyResolver(store_ref)

            new_args = list(args)
            new_kwargs = kwargs.copy()

            # Resolve Kwargs
            for k, v in new_kwargs.items():
                if isinstance(v, str) and "{" in v and "}" in v:
                    try:
                        resolved = await resolver.execute_dag(v, entity_id)
                        new_kwargs[k] = resolved
                    except Exception as e:
                        logger.warning(
                            f"DAG resolution failed for kwarg {k}", error=str(e)
                        )

            # Resolve Args (skip for now to avoid positional confusion, or iterate)
            # Users defined @retriever usually use kwargs for clear query inputs like `query="..."`

            return tuple(new_args), new_kwargs

        # Async Wrapper Support
        import inspect

        if inspect.iscoroutinefunction(func):

            @functools.wraps(func)
            async def async_wrapper(*args: Any, **kwargs: Any) -> List[Dict[str, Any]]:
                # 1. Resolve DAG Dependencies
                args, kwargs = await _resolve_args(args, kwargs)

                store_backend = getattr(ret_obj, "_cache_backend", None)
                if ret_obj.cache_ttl and store_backend:
                    try:
                        import json
                        import hashlib

                        key_parts = [r_name, str(args), str(kwargs)]
                        key_str = json.dumps(key_parts, sort_keys=True, default=str)
                        key_hash = hashlib.sha256(key_str.encode("utf-8")).hexdigest()
                        cache_key = f"meridian:retriever:{r_name}:{key_hash}"

                        # Try fetch
                        cached = await store_backend.get(cache_key)
                        if cached:
                            logger.info(f"Retriever Cache Hit: {r_name}")
                            return cast(List[Dict[str, Any]], json.loads(cached))

                        # Miss -> Call
                        result = await func(*args, **kwargs)

                        # Store
                        ttl_sec = int(ret_obj.cache_ttl.total_seconds())
                        await store_backend.set(
                            cache_key, json.dumps(result), ex=ttl_sec
                        )
                        return cast(List[Dict[str, Any]], result)

                    except Exception as e:
                        logger.warning(f"Retriever Caching Error: {e}")

                return cast(List[Dict[str, Any]], await func(*args, **kwargs))

            return async_wrapper  # type: ignore[return-value]

        # Sync Wrapper
        @functools.wraps(func)
        def sync_wrapper(*args: Any, **kwargs: Any) -> List[Dict[str, Any]]:
            # Sync cannot await DAG resolution simply.
            # Limitation: DAG Wiring only supported for Async Retrievers in V1.

            has_template = any(isinstance(v, str) and "{" in v for v in kwargs.values())
            if has_template:
                logger.warning(
                    "sync_retriever_dag_skipped",
                    reason="DAG resolution requires async retriever",
                )

            result = func(*args, **kwargs)

            import inspect

            if inspect.isawaitable(result):
                # Safety check: if user wrapped an async func but it wasn't detected as coroutinefunction
                # (e.g. partial or other callable), we can't await it here.
                raise RuntimeError(
                    "Sync retriever returned awaitable. Use async def for async retrievers."
                )

            return cast(List[Dict[str, Any]], list(result))

        return sync_wrapper

    return decorator
