import inspect
import functools
import contextvars
from typing import Callable
from langfuse import Langfuse


class ContextManager:
    context = {
        "langfuse": contextvars.ContextVar("langfuse"),
        "stack": contextvars.ContextVar("stack"),
        "observation": contextvars.ContextVar("observation"),
    }

    def __getattr__(self, name):
        if name not in self.context:
            raise AttributeError(f"'ContextManager' object has no attribute '{name}'")
        return self.context[name].get(None)

    def __setattr__(self, name, value):
        if name not in self.context:
            raise AttributeError(f"'ContextManager' object has no attribute '{name}'")
        return self.context[name].set(value)

def default_input_fn(args, kwargs):
    return {**kwargs, "*args": args} if args else kwargs

def ismethod(func: Callable) -> bool:
    return 'self' in inspect.signature(func).parameters

def trace(
    input_fn: Callable = default_input_fn,
    output_fn: Callable = lambda x: x,
    **decorator_kwargs,
):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            ctx = ContextManager()
            if not ctx.langfuse:
                ctx.langfuse = Langfuse()

            # init stack
            input_args = args[1:] if ismethod(func) else args
            ctx.stack = [
                ctx.langfuse.trace(**decorator_kwargs, input=input_fn(input_args, kwargs)),
            ]

            # run target
            output = func(*args, **kwargs)
            ctx.stack[-1].update(output=output_fn(output))
            ctx.langfuse.flush()
            return output
        
        # FIXME: it's just a copy of wrapper and await func(...)
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            ctx = ContextManager()
            if not ctx.langfuse:
                ctx.langfuse = Langfuse()

            # init stack
            input_args = args[1:] if ismethod(func) else args
            ctx.stack = [
                ctx.langfuse.trace(**decorator_kwargs, input=input_fn(input_args, kwargs)),
            ]

            # run target
            output = await func(*args, **kwargs)
            ctx.stack[-1].update(output=output_fn(output))
            ctx.langfuse.flush()
            return output

        return async_wrapper if inspect.iscoroutinefunction(func) else wrapper
    return decorator

def span(
    input_fn: Callable = default_input_fn,
    output_fn: Callable = lambda x: x,
    **decorator_kwargs,
):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            ctx = ContextManager()
            input_args = args[1:] if ismethod(func) else args
            ctx.stack.append(
                ctx.stack[-1].span(**decorator_kwargs, input=input_fn(input_args, kwargs))
            )

            try:
                output = func(*args, **kwargs)
                ctx.stack[-1].end(output=output_fn(output))
                return output
            finally:
                ctx.observation = ctx.stack.pop()

        # FIXME: it's just a copy of wrapper and await func(...)
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            ctx = ContextManager()
            input_args = args[1:] if ismethod(func) else args
            ctx.stack.append(
                ctx.stack[-1].span(**decorator_kwargs, input=input_fn(input_args, kwargs))
            )

            try:
                output = await func(*args, **kwargs)
                ctx.stack[-1].end(output=output_fn(output))
                return output
            finally:
                ctx.observation = ctx.stack.pop()

        return async_wrapper if inspect.iscoroutinefunction(func) else wrapper
    return decorator

def generation(
    input_fn: Callable = default_input_fn,
    output_fn: Callable = lambda x: x,
    usage_fn: Callable = None,
    **decorator_kwargs,
):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            ctx = ContextManager()
            input_args = args[1:] if ismethod(func) else args
            ctx.stack.append(
                ctx.stack[-1].generation(**decorator_kwargs, input=input_fn(input_args, kwargs))
            )

            try:
                output = func(*args, **kwargs)
                if usage_fn:
                    ctx.stack[-1].end(output=output_fn(output), usage=usage_fn(output))
                else:
                    ctx.stack[-1].end(output=output_fn(output))
                return output
            finally:
                ctx.observation = ctx.stack.pop()

        # FIXME: it's just a copy of wrapper and await func(...)
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            ctx = ContextManager()
            input_args = args[1:] if ismethod(func) else args
            ctx.stack.append(
                ctx.stack[-1].generation(**decorator_kwargs, input=input_fn(input_args, kwargs))
            )

            try:
                output = await func(*args, **kwargs)
                if usage_fn:
                    ctx.stack[-1].end(output=output_fn(output), usage=usage_fn(output))
                else:
                    ctx.stack[-1].end(output=output_fn(output))
                return output
            finally:
                ctx.observation = ctx.stack.pop()

        return async_wrapper if inspect.iscoroutinefunction(func) else wrapper
    return decorator

def event(
    input_fn: Callable = default_input_fn,
    output_fn: Callable = lambda x: x,
    **decorator_kwargs,
):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            ctx = ContextManager()
            output = func(*args, **kwargs)
            input_args = args[1:] if ismethod(func) else args
            ctx.observation = ctx.stack[-1].event(
                **decorator_kwargs, 
                input=input_fn(input_args, kwargs), 
                output=output_fn(output),
            )
            return output

        # FIXME: it's just a copy of wrapper and await func(...)
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            ctx = ContextManager()
            output = await func(*args, **kwargs)
            input_args = args[1:] if ismethod(func) else args
            ctx.observation = ctx.stack[-1].event(
                **decorator_kwargs, 
                input=input_fn(input_args, kwargs), 
                output=output_fn(output),
            )
            return output

        return async_wrapper if inspect.iscoroutinefunction(func) else wrapper
    return decorator

def current_observation():
    return ContextManager().stack[-1]

def previous_observation():
    return ContextManager().observation