from contextvars import ContextVar
from typing import Optional

from yera.agents.context import get_agent_context
from yera.events import input_echo, markdown, system_prompt
from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel
from yera.models.llm_workspace import LLMWorkspace

_current_llm_context: ContextVar[Optional["LLMContext"]] = ContextVar(
    "llm_context", default=None
)


class LLMContext:
    def __init__(
        self,
        *,
        interface: BaseLLMInterface,
    ):
        self.workspace = None
        self.interface = interface
        self.agent_meta = None
        self._token = None

    def _push_event(self, role, content):
        if role == "user":
            input_echo(content)
        elif role == "system":
            system_prompt(content)
        elif role == "agent":
            markdown(content)
        else:
            raise ValueError(f"Unknown role: {role}")

    def add_sys_line(self, content):
        self._push_event("system", content)
        self.workspace.add_sys_message(content)

    def add_user_line(self, content):
        self._push_event("user", content)
        self.workspace.add_user_message(content)

    def add_assistant_line(self, content):
        self.workspace.add_assistant_message(content)

    def set(self, key, value):
        self.workspace.variables[key] = value

    def get(self, key, default=None):
        return self.workspace.variables.get(key, default)

    def __getitem__(self, key):
        return self.workspace.variables[key]

    def prompt_chat(self, prompt: str, **kwargs) -> str:
        self.add_user_line(prompt)
        stream = self.interface.chat(self.workspace.messages, **kwargs)

        tokens = []
        # Use markdown as context manager to create a single stream
        # All tokens will go to the same block with sequential chunk_ids
        with markdown() as md:
            for token in stream:
                md.append(token)
                tokens.append(token)

        response = "".join(tokens)
        self.add_assistant_line(response)
        return response

    def prompt_struct(self, prompt: str, cls: type[TBaseModel], **kwargs) -> TBaseModel:
        self.add_user_line(prompt)
        stream = self.interface.make_struct(self.workspace.messages, cls, **kwargs)

        tokens = []
        # Use markdown as context manager to create a single stream
        # All tokens will go to the same block with sequential chunk_ids
        with markdown() as md:
            for token in stream:
                md.append(token)
                tokens.append(token)

        response_json = "".join(tokens)
        self.add_assistant_line(response_json)
        return cls.model_validate_json(response_json)

    def _handle_workspace(self) -> LLMWorkspace:
        exterior_model = llm_context(error_when_absent=False)
        if exterior_model and exterior_model.agent_meta == self.agent_meta:
            return exterior_model.workspace
        return LLMWorkspace()

    def __enter__(self):
        self.agent_meta = get_agent_context().metadata
        self.workspace = self._handle_workspace()
        self.interface.start()
        self._token = _current_llm_context.set(self)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.interface.stop()
        _current_llm_context.reset(self._token)


def llm_context(error_when_absent: bool = True) -> LLMContext:
    llm = _current_llm_context.get()
    if llm is None and error_when_absent:
        raise RuntimeError("No current LLM")
    return llm
