from .collector import ThreadManager
from .client import Laminar
from .providers import Provider
from .providers.fallback import FallbackProvider
from .tracing_types import EvaluateEvent, Event, Span, Trace
from .types import PipelineRunResponse
from .utils import PROVIDER_NAME_TO_OBJECT, is_async_iterator, is_iterator

from contextvars import ContextVar
from typing import Any, AsyncGenerator, Generator, Literal, Optional, Union
import atexit
import datetime
import dotenv
import inspect
import logging
import os
import pydantic
import uuid


_lmnr_stack_context: ContextVar[list[Union[Span, Trace]]] = ContextVar(
    "lmnr_stack_context", default=[]
)
_root_trace_id_context: ContextVar[Optional[str]] = ContextVar(
    "root_trace_id_context", default=None
)


class LaminarContextManager:
    _log = logging.getLogger("laminar.context_manager")

    def __init__(
        self,
        project_api_key: str = None,
        threads: int = 1,
        max_task_queue_size: int = 1000,
    ):
        self.project_api_key = project_api_key or os.environ.get("LMNR_PROJECT_API_KEY")
        if not self.project_api_key:
            dotenv_path = dotenv.find_dotenv(usecwd=True)
            self.project_api_key = dotenv.get_key(
                dotenv_path=dotenv_path, key_to_get="LMNR_PROJECT_API_KEY"
            )
        self.laminar = Laminar(project_api_key=self.project_api_key)
        self.thread_manager = ThreadManager(
            client=self.laminar,
            max_task_queue_size=max_task_queue_size,
            threads=threads,
        )
        # atexit executes functions last in first out, so we want to make sure
        # that we finalize the trace before thread manager is closed, so the updated
        # trace is sent to the server
        atexit.register(self._force_finalize_trace)

    def observe_start(
        self,
        # span attributes
        name: str,
        input: Optional[Any] = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
        span_type: Literal["DEFAULT", "LLM"] = "DEFAULT",
        check_event_names: list[str] = None,
        # trace attributes
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        release: Optional[str] = None,
    ) -> Span:
        trace_id = _root_trace_id_context.get()
        if not trace_id:
            session_id = session_id or str(uuid.uuid4())
            trace_id = uuid.uuid4()
            trace = self.update_trace(
                id=trace_id,
                user_id=user_id,
                session_id=session_id,
                release=release,
            )
            _root_trace_id_context.set(trace.id)
            _lmnr_stack_context.set([trace])

        parent = _lmnr_stack_context.get()[-1] if _lmnr_stack_context.get() else None
        parent_span_id = parent.id if isinstance(parent, Span) else None
        span = self.create_span(
            name=name,
            trace_id=trace_id,
            input=input,
            metadata=metadata,
            attributes=attributes,
            parent_span_id=parent_span_id,
            span_type=span_type,
            check_event_names=check_event_names,
        )
        stack = _lmnr_stack_context.get()
        _lmnr_stack_context.set(stack + [span])
        return span

    def observe_end(
        self,
        span: Span,
        provider_name: str = None,
        result: Optional[Any] = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
        error: Optional[Exception] = None,
    ) -> Any:
        stack = _lmnr_stack_context.get()
        if not stack:
            return
        provider = PROVIDER_NAME_TO_OBJECT.get(provider_name, FallbackProvider())
        new_stack = stack[:-1]
        _lmnr_stack_context.set(new_stack)

        if len(new_stack) == 1 and isinstance(stack[0], Trace):
            trace = stack[0]
            self.update_trace(
                id=trace.id,
                user_id=trace.userId,
                session_id=trace.sessionId,
                release=trace.release,
                metadata=metadata,
            )
            _root_trace_id_context.set(None)
            _lmnr_stack_context.set([])

        if error is not None:
            self.update_current_trace(success=False)

        if inspect.isgenerator(result) or is_iterator(result):
            return self._collect_generator_result(
                provider=provider,
                generator=result,
                span=span,
                metadata=metadata,
                attributes=attributes,
            )
        elif inspect.isasyncgen(result) or is_async_iterator(result):
            return self._collect_async_generator_result(
                provider=provider,
                generator=result,
                span=span,
                metadata=metadata,
                attributes=attributes,
            )
        if span.spanType == "LLM" and error is None:
            attributes = self._extract_llm_attributes_from_response(
                provider=provider, response=result
            )
        return self._finalize_span(
            span,
            provider=provider,
            result=error or result,
            metadata=metadata,
            attributes=attributes,
        )

    def update_current_span(
        self,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
        evaluate_events: list[EvaluateEvent] = None,
        override: bool = False,
    ):
        stack = _lmnr_stack_context.get()
        if not stack:
            return
        span = stack[-1]
        new_metadata = (
            metadata if override else {**(span.metadata or {}), **(metadata or {})}
        )
        new_evaluate_events = (
            evaluate_events
            if override
            else span.evaluateEvents + (evaluate_events or [])
        )
        new_attributes = (
            attributes
            if override
            else {**(span.attributes or {}), **(attributes or {})}
        )
        self.update_span(
            span=span,
            metadata=new_metadata,
            evaluate_events=new_evaluate_events,
            attributes=new_attributes,
        )

    def update_current_trace(
        self,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        release: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
        success: bool = True,
    ):
        existing_trace = (
            _lmnr_stack_context.get()[0] if _lmnr_stack_context.get() else None
        )
        if not existing_trace:
            return
        self.update_trace(
            id=existing_trace.id,
            user_id=user_id or existing_trace.userId,
            session_id=session_id or existing_trace.sessionId,
            release=release or existing_trace.release,
            metadata=metadata or existing_trace.metadata,
            success=success if success is not None else existing_trace.success,
        )

    def update_trace(
        self,
        id: uuid.UUID,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        release: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
        success: bool = True,
    ) -> Trace:
        trace = Trace(
            id=id,
            user_id=user_id,
            session_id=session_id,
            release=release,
            metadata=metadata,
            success=success,
        )
        self._add_observation(trace)
        return trace

    def create_span(
        self,
        name: str,
        trace_id: uuid.UUID,
        start_time: Optional[datetime.datetime] = None,
        span_type: Literal["DEFAULT", "LLM"] = "DEFAULT",
        id: Optional[uuid.UUID] = None,
        parent_span_id: Optional[uuid.UUID] = None,
        input: Optional[Any] = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
        check_event_names: list[str] = None,
    ) -> Span:
        """Internal method to create a span object. Use `ObservationContext.span` instead."""
        span = Span(
            name=name,
            trace_id=trace_id,
            start_time=start_time or datetime.datetime.now(datetime.timezone.utc),
            id=id,
            parent_span_id=parent_span_id,
            input=input,
            metadata=metadata,
            attributes=attributes,
            span_type=span_type,
            evaluate_events=check_event_names or [],
        )
        return span

    def update_span(
        self,
        span: Span,
        finalize: bool = False,
        input: Optional[Any] = None,
        end_time: Optional[datetime.datetime] = None,
        output: Optional[Any] = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
        evaluate_events: Optional[list[EvaluateEvent]] = None,
        override: bool = False,
    ) -> Span:
        """Internal method to update a span object. Use `SpanContext.update()` instead."""
        span.update(
            input=input or span.input,
            output=output or span.output,
            end_time=end_time,
            metadata=metadata,
            attributes=attributes,
            evaluate_events=evaluate_events,
            override=override,
        )
        if finalize:
            self._add_observation(span)
        return span

    def event(
        self,
        name: str,
        value: Optional[Union[str, int, float, bool]] = None,
        timestamp: Optional[datetime.datetime] = None,
    ):
        span = _lmnr_stack_context.get()[-1] if _lmnr_stack_context.get() else None
        if not span or not isinstance(span, Span):
            self._log.warning(f"No active span to send event. Ignoring event. {name}")
            return
        event = Event(
            name=name,
            span_id=span.id,
            timestamp=timestamp,
            value=value,
        )
        span.add_event(event)

    def evaluate_event(self, name: str, data: str):
        stack = _lmnr_stack_context.get()
        if not stack or not isinstance(stack[-1], Span):
            self._log.warning(
                f"No active span to add check event. Ignoring event. {name}"
            )
            return
        stack[-1].evaluateEvents.append(
            EvaluateEvent(
                name=name,
                data=data,
                timestamp=datetime.datetime.now(datetime.timezone.utc),
            )
        )

    def run_pipeline(
        self,
        pipeline: str,
        inputs: dict[str, Any],
        env: dict[str, str] = {},
        metadata: dict[str, str] = {},
    ) -> PipelineRunResponse:
        span = _lmnr_stack_context.get()[-1] if _lmnr_stack_context.get() else None
        span_id = span.id if isinstance(span, Span) else None
        trace = _lmnr_stack_context.get()[0] if _lmnr_stack_context.get() else None
        trace_id = trace.id if isinstance(trace, Trace) else None
        return self.laminar.run(
            pipeline=pipeline,
            inputs=inputs,
            env=env,
            metadata=metadata,
            parent_span_id=span_id,
            trace_id=trace_id,
        )

    def _force_finalize_trace(self):
        # TODO: flush in progress spans as error?
        pass

    def _add_observation(self, observation: Union[Span, Trace]) -> bool:
        return self.thread_manager.add_task(observation)

    def _extract_llm_attributes_from_response(
        self,
        provider: Provider,
        response: Union[str, dict[str, Any], pydantic.BaseModel],
    ) -> dict[str, Any]:
        return provider.extract_llm_attributes_from_response(response)

    def _stream_list_to_dict(
        self, provider: Provider, response: list[Any]
    ) -> dict[str, Any]:
        return provider.stream_list_to_dict(response)

    def _extract_llm_output(
        self,
        provider: Provider,
        result: Union[dict[str, Any], pydantic.BaseModel],
    ) -> str:
        return provider.extract_llm_output(result)

    def _finalize_span(
        self,
        span: Span,
        provider: Provider = None,
        result: Optional[Any] = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
    ) -> Any:
        self.update_span(
            span=span,
            finalize=True,
            output=(
                result
                if span.spanType != "LLM"
                else self._extract_llm_output(provider, result)
            ),
            metadata=metadata,
            attributes=attributes,
        )
        return result

    def _collect_generator_result(
        self,
        generator: Generator,
        span: Span,
        provider: Provider = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
    ) -> Generator:
        items = []
        try:
            for item in generator:
                items.append(item)
                yield item

        finally:
            output = items
            if all(isinstance(item, str) for item in items):
                output = "".join(items)
            if span.spanType == "LLM":
                collected = self._stream_list_to_dict(
                    provider=provider, response=output
                )
                attributes = self._extract_llm_attributes_from_response(
                    provider=provider, response=collected
                )
            self._finalize_span(
                span=span,
                provider=provider,
                result=collected,
                metadata=metadata,
                attributes=attributes,
            )

    async def _collect_async_generator_result(
        self,
        generator: AsyncGenerator,
        span: Span,
        provider: Provider = None,
        metadata: Optional[dict[str, Any]] = None,
        attributes: Optional[dict[str, Any]] = None,
    ) -> AsyncGenerator:
        items = []
        try:
            async for item in generator:
                items.append(item)
                yield item

        finally:
            output = items
            if all(isinstance(item, str) for item in items):
                output = "".join(items)
            if span.spanType == "LLM":
                collected = self._stream_list_to_dict(
                    provider=provider, response=output
                )
                attributes = self._extract_llm_attributes_from_response(
                    provider=provider, response=collected
                )
            self._finalize_span(
                span=span,
                provider=provider,
                result=collected,
                metadata=metadata,
                attributes=attributes,
            )


# TODO: add lock for thread safety
class LaminarSingleton:
    _instance = None
    _l: Optional[LaminarContextManager] = None

    def __new__(cls):
        if not cls._instance:
            cls._instance = super(LaminarSingleton, cls).__new__(cls)
        return cls._instance

    def get(cls, *args, **kwargs) -> LaminarContextManager:
        if not cls._l:
            cls._l = LaminarContextManager(*args, **kwargs)
        return cls._l
