import dspy
import os
import httpx
from twopir.llms.base import (
    ModelEntry,
    ChatMessages,
    LLM,
)

COSTS = {
    "amazon.nova-lite-v1:0": {
        "per_input_token_cost": 0.06 / 1000000.0,
        "per_output_token_cost": 0.24 / 1000000.0,
    },
}


class BedrockLlm(LLM):
    """Bedrock LLM.
    This class requires AWS_BEARER_TOKEN_BEDROCK to be in the environment. It will call models hosted there.
    """

    def get_id(self) -> str:
        return f"bedrock:{self.name}"

    def __init__(self, name: str, client: httpx.AsyncClient) -> None:
        self.name = name
        self._client = client
        self._per_input_token_cost = COSTS[self.name]["per_input_token_cost"]
        self._per_output_token_cost = COSTS[self.name]["per_output_token_cost"]

    def _to_converse_format(self, messages: ChatMessages) -> list:
        system = None
        converted_messages = []
        for message in messages:
            if message["role"] == "system":
                system = [{"text": message["content"]}]
            else:
                converted_messages.append(
                    {
                        "role": message["role"],
                        "content": [{"text": message["content"]}],  # type: ignore
                    }
                )
        return system, converted_messages

    async def do_generate(
        self,
        messages: ChatMessages,
        **kwargs,
    ) -> str:
        system, converted_messages = self._to_converse_format(messages)
        json_data = {
            "messages": converted_messages,
            "inferenceConfig": {
                "maxTokens": self.max_tokens,
                "temperature": self.temperature,
            },
        }
        if system:
            json_data["system"] = system
        response = await self._client.post(
            f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{self.name}/converse",
            json=json_data,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {os.environ['AWS_BEARER_TOKEN_BEDROCK']}",
            },
        )
        response.raise_for_status()
        return response.json()["output"]["message"]["content"][0]["text"]

    def dspy_lm(self) -> dspy.LM:
        raise ValueError("Not implemented yet!")


NOVA_LITE = ModelEntry(
    name="nova-lite",
    get_with_client=lambda client: BedrockLlm(
        name="amazon.nova-lite-v1:0",
        client=client,
    ),
)