from typing import Optional

from langchain_openai import ChatOpenAI
from pydantic import SecretStr

from openkosmos_ai.common.model import AIBaseConfig
from openkosmos_ai.nodes.base_node import BaseStateNode


class OpenAIClientConfig(AIBaseConfig):
    base_url: str
    api_key: str
    model: Optional[str] = None
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.8
    presence_penalty: Optional[float] = 1.03


class OpenAIClientNode(BaseStateNode[OpenAIClientConfig]):
    def __init__(self, config: OpenAIClientConfig, tools=None):
        super().__init__(config)
        self.llm = ChatOpenAI(
            model=config.model,
            base_url=config.base_url,
            api_key=SecretStr(config.api_key),
            temperature=config.temperature,
            top_p=config.top_p,
            presence_penalty=config.presence_penalty
        )
        if tools is not None:
            self.llm = self.llm.bind_tools(tools)

    def model(self):
        return self.llm

    def invoke(self, messages, output_schema=None):
        if output_schema is None:
            return self.llm.invoke(messages)
        else:
            return self.llm.with_structured_output(output_schema).invoke(messages)

    def chat(self, messages):
        response = self.llm.invoke(messages)
        return response.content_blocks
