from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Dict, Optional, Union
from uuid import UUID, uuid4

from freeplay.errors import FreeplayClientError
from freeplay.model import JSONValue, SpanKind, TestRunInfo
from freeplay.support import CallSupport, CustomMetadata


@dataclass
class SessionInfo:
    session_id: str
    custom_metadata: CustomMetadata


class TraceInfo:
    session_id: str
    trace_id: str
    input: Optional[JSONValue] = None
    agent_name: Optional[str] = None
    kind: Optional[SpanKind] = None
    name: Optional[str] = None
    parent_id: Optional[UUID] = None
    custom_metadata: CustomMetadata = None
    start_time: datetime
    _call_support: CallSupport

    def __init__(
        self,
        trace_id: str,
        session_id: str,
        _call_support: CallSupport,
        input: Optional[JSONValue] = None,
        agent_name: Optional[str] = None,
        parent_id: Optional[UUID] = None,
        custom_metadata: CustomMetadata = None,
        kind: Optional[SpanKind] = None,
        name: Optional[str] = None,
        start_time: Optional[datetime] = None,
    ):
        self.trace_id = trace_id
        self.session_id = session_id
        self.input = input
        self.agent_name = agent_name
        self.parent_id = parent_id
        self.custom_metadata = custom_metadata
        self._call_support = _call_support
        self.kind = kind
        self.name = name
        self.start_time = start_time or datetime.now(timezone.utc)

    def record_output(
        self,
        project_id: str,
        output: JSONValue,
        eval_results: Optional[Dict[str, Union[bool, float]]] = None,
        test_run_info: Optional[TestRunInfo] = None,
        end_time: Optional[datetime] = None,
    ) -> None:
        if self.input is None:
            raise FreeplayClientError("Input must be set before recording output")
        self._call_support.record_trace(
            project_id,
            self.session_id,
            self.trace_id,
            self.input,
            output,
            agent_name=self.agent_name,
            parent_id=self.parent_id,
            custom_metadata=self.custom_metadata,
            eval_results=eval_results,
            test_run_info=test_run_info,
            kind=self.kind,
            name=self.name,
            start_time=self.start_time,
            end_time=end_time or datetime.now(timezone.utc),
        )


@dataclass
class Session:
    session_id: str
    custom_metadata: CustomMetadata

    def __init__(
        self,
        session_id: str,
        custom_metadata: CustomMetadata,
        _call_support: CallSupport,
    ):
        self.session_id = session_id
        self.custom_metadata = custom_metadata
        self._session_info = SessionInfo(self.session_id, self.custom_metadata)
        self._call_support = _call_support

    @property
    def session_info(self) -> SessionInfo:
        return self._session_info

    def create_trace(
        self,
        input: JSONValue,
        agent_name: Optional[str] = None,
        parent_id: Optional[UUID] = None,
        custom_metadata: CustomMetadata = None,
        kind: Optional[SpanKind] = None,
        name: Optional[str] = None,
    ) -> TraceInfo:
        return TraceInfo(
            trace_id=str(uuid4()),
            session_id=self.session_id,
            parent_id=parent_id,
            input=input,
            agent_name=agent_name,
            custom_metadata=custom_metadata,
            _call_support=self._call_support,
            kind=kind,
            name=name,
        )

    def restore_trace(
        self,
        trace_id: UUID,
        input: Optional[JSONValue],
        agent_name: Optional[str] = None,
        parent_id: Optional[UUID] = None,
        custom_metadata: CustomMetadata = None,
        kind: Optional[SpanKind] = None,
        name: Optional[str] = None,
    ) -> TraceInfo:
        return TraceInfo(
            trace_id=str(trace_id),
            session_id=self.session_id,
            input=input,
            agent_name=agent_name,
            parent_id=parent_id,
            custom_metadata=custom_metadata,
            _call_support=self._call_support,
            kind=kind,
            name=name,
        )


class Sessions:
    def __init__(self, call_support: CallSupport):
        self.call_support = call_support

    def create(self, custom_metadata: CustomMetadata = None) -> Session:
        return Session(
            session_id=str(uuid4()),
            custom_metadata=custom_metadata,
            _call_support=self.call_support,
        )

    def delete(self, project_id: str, session_id: str) -> None:
        self.call_support.delete_session(project_id, session_id)

    def restore_session(
        self, session_id: str, custom_metadata: CustomMetadata = None
    ) -> Session:
        return Session(
            session_id=session_id,
            custom_metadata=custom_metadata,
            _call_support=self.call_support,
        )
