# LLM utils

import json

from openai import OpenAI
from os import environ


def tool(name, description, parameters):
    def decorator(fn):
        fn.tool_name = name
        fn.tool_definition = {
            "type": "function",
            "name": name,
            "description": description,
            "parameters": parameters,
        }
        return fn
    return decorator


class LLM:
    def __init__(self, api_key=None, base_url=None, model=None, client=None, **kwargs):
        self.api_key = api_key or environ.get("LLM_API_KEY")
        if not self.api_key:
            raise ValueError("No LLM API key found")
        self.base_url = base_url or environ.get("LLM_BASE_URL") or "https://px.axess-ai.com/v1"
        self.model = model
        self.client = client or OpenAI(base_url=self.base_url, api_key=self.api_key, **kwargs)

    def with_model(self, model):
        return LLM(api_key=self.api_key, base_url=self.base_url, model=model, client=self.client)

    def complete(self, user_prompt, system_prompt=None, model=None, **kwargs):
        use_model = model or self.model
        if not use_model:
            raise ValueError("Model not provided in constructor nor at completion")
        response = self.client.responses.create(
            model=use_model,
            input=user_prompt,
            **({"instructions": system_prompt} if system_prompt else {}),
            **kwargs,
        )
        return response.output_text

    def json_complete(self, user_prompt, response_format, system_prompt=None, model=None, **kwargs):
        use_model = model or self.model
        if not use_model:
            raise ValueError("Model not provided in constructor nor at completion")
        response = self.client.responses.parse(
            model=use_model,
            input=user_prompt,
            text_format=response_format,
            **({"instructions": system_prompt} if system_prompt else {}),
            **kwargs,
        )
        return response.output_parsed.model_dump()

    def tool_complete(self, user_prompt, tools, system_prompt=None, model=None, max_rounds=10, **kwargs):
        use_model = model or self.model
        if not use_model:
            raise ValueError("Model not provided in constructor nor at completion")
        tool_map = {t.tool_name: t for t in tools}
        tool_definitions = [t.tool_definition for t in tools]
        input = [{"role": "user", "content": user_prompt}]
        for _ in range(max_rounds):
            response = self.client.responses.create(
                model=use_model,
                input=input,
                tools=tool_definitions,
                **({"instructions": system_prompt} if system_prompt else {}),
                **kwargs,
            )
            function_calls = [item for item in response.output if item.type == "function_call"]
            if not function_calls:
                return response.output_text
            input.extend(response.output)
            for call in function_calls:
                result = tool_map[call.name](**json.loads(call.arguments))
                input.append({
                    "type": "function_call_output",
                    "call_id": call.call_id,
                    "output": result if isinstance(result, str) else json.dumps(result),
                })
        raise ValueError(f"Tool calling did not complete within {max_rounds} rounds")

    def embed(self, content, model=None, **kwargs):
        use_model = model or self.model
        if not use_model:
            raise ValueError("Model not provided in constructor nor at completion")
        embedding = self.client.embeddings.create(model=use_model, input=content, encoding_format="float", **kwargs)
        return embedding.data[0].embedding
