import asyncio
import json
from contextvars import ContextVar
from functools import wraps
from typing import Callable, Dict, Optional
from uuid import uuid4

from ..logger.components.retrieval import Retrieval, RetrievalConfig
from ..logger.logger import Logger
from .span import current_span
from .trace import current_logger, current_trace

_retrieval_ctx_var: ContextVar[Optional[Retrieval]] = ContextVar(
    "maxim_ctx_retrieval", default=None
)


def current_retrieval() -> Optional[Retrieval]:
    return _retrieval_ctx_var.get()


def retrieval(
    logger: Optional[Logger] = None,
    id: Optional[str] or Optional[Callable] = None,
    input: Optional[str] or Optional[Callable] = None,
    name: Optional[str] = None,
    tags: Optional[Dict[str, str]] = None,
):
    def decorator(func):        
        if asyncio.iscoroutinefunction(func):
            @wraps(func)
            async def wrapper(*args, **kwargs):
                # First check if the logger is available
                maxim_logger = logger
                if maxim_logger is None:
                    if current_logger() is None:
                        raise ValueError(
                            "[MaximSDK]: no logger found. either call this function from a @trace decorated function or pass a logger"
                        )
                    maxim_logger = current_logger()
                # Here there should be an active span or active trace
                # If none of this is the case then we raise an error
                if current_span() is None and current_trace() is None:
                    if maxim_logger.raise_exceptions:
                        raise ValueError(
                            "[MaximSDK]: no trace or span found. either call this function from a @trace or @span decorated function"
                        )
                    else:
                        logging.warning(
                            "[MaximSDK]: no trace or span found. either call this function from a @trace or @span decorated function"
                        )
                actual_retrieval_id = id() if callable(id) else id
                if actual_retrieval_id is None:
                    actual_retrieval_id = str(uuid4())
                config = RetrievalConfig(id=actual_retrieval_id, name=name, tags=tags)
                retrieval: Retrieval = None
                if current_span() is not None:
                    retrieval = current_span().retrieval(config)
                elif current_trace() is not None:
                    retrieval = current_trace().retrieval(config)
                # Checking for input
                actual_input = input() if callable(input) else input
                if actual_input is not None:
                    retrieval.input(actual_input)
                token = _retrieval_ctx_var.set(retrieval)
                try:
                    result = await func(*args, **kwargs)
                    # Here we will check if retrieval has output
                    if not retrieval.is_output_set:
                        if isinstance(result, (dict)):
                            retrieval.output(json.dumps(result))
                        elif isinstance(result, list):
                            retrieval.output([json.dumps(item) for item in result])
                        elif isinstance(result, str):
                            retrieval.output(result)
                    return result
                finally:
                    retrieval.end()
                    _retrieval_ctx_var.reset(token)
        else:
            @wraps(func)
            def wrapper(*args, **kwargs):
                # First check if the logger is available
                maxim_logger = logger
                if maxim_logger is None:
                    if current_logger() is None:
                        raise ValueError(
                            "[MaximSDK]: no logger found. either call this function from a @trace decorated function or pass a logger"
                        )
                    maxim_logger = current_logger()
                # Here there should be an active span or active trace
                # If none of this is the case then we raise an error
                if current_span() is None and current_trace() is None:
                    if maxim_logger.raise_exceptions:
                        raise ValueError(
                            "[MaximSDK]: no trace or span found. either call this function from a @trace or @span decorated function"
                        )
                    else:
                        logging.warning(
                            "[MaximSDK]: no trace or span found. either call this function from a @trace or @span decorated function"
                        )
                actual_retrieval_id = id() if callable(id) else id
                if actual_retrieval_id is None:
                    actual_retrieval_id = str(uuid4())
                config = RetrievalConfig(id=actual_retrieval_id, name=name, tags=tags)
                retrieval: Retrieval = None
                if current_span() is not None:
                    retrieval = current_span().retrieval(config)
                elif current_trace() is not None:
                    retrieval = current_trace().retrieval(config)
                # Checking for input
                actual_input = input() if callable(input) else input
                if actual_input is not None:
                    retrieval.input(actual_input)
                token = _retrieval_ctx_var.set(retrieval)
                try:
                    result = func(*args, **kwargs)
                    # Here we will check if retrieval has output
                    if not retrieval.is_output_set:
                        if isinstance(result, (dict)):
                            retrieval.output(json.dumps(result))
                        elif isinstance(result, list):
                            retrieval.output([json.dumps(item) for item in result])
                        elif isinstance(result, str):
                            retrieval.output(result)
                    return result
                finally:
                    retrieval.end()
                    _retrieval_ctx_var.reset(token)
        return wrapper

    return decorator
