import json
import boto3
from analogai.deepthink.infrastructure.gateways.llm.gateway import LlmGateway
from analogai.deepthink.infrastructure.gateways.llm import config


class BedrockLlmGatewayImpl(LlmGateway):
    def __init__(self, system_prompt: str | None = None, max_tokens: int = 4096, model: str | None = None):
        if not config.AWS_ACCESS_KEY_ID or not config.AWS_SECRET_ACCESS_KEY:
            raise ValueError("AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY must be set for Bedrock gateway")
        if not model:
            raise ValueError("Model name must be provided for Bedrock gateway")
        self.model = model
        self.default_system_prompt = system_prompt or ""
        self.default_max_tokens = min(max_tokens, 4096)
        self.client = boto3.client(
            service_name="bedrock-runtime",
            region_name="us-east-1",
            aws_access_key_id=config.AWS_ACCESS_KEY_ID,
            aws_secret_access_key=config.AWS_SECRET_ACCESS_KEY,
        )

    def set_model(self, model_id: str) -> None:
        self.model = model_id

    def set_system_prompt(self, system_prompt: str) -> None:
        self.default_system_prompt = system_prompt

    def generate_completion(
        self,
        user_message: str,
    ) -> str:
        final_system_prompt = self.default_system_prompt
        max_tokens = self.default_max_tokens

        request_body = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": max_tokens,
            "messages": [
                {"role": "user", "content": user_message}
            ],
            "system": final_system_prompt,
        }

        response = self.client.invoke_model(
            modelId=self.model,
            body=json.dumps(request_body),
        )
        response_body = json.loads(response.get("body").read())
        return response_body["content"][0]["text"]

    async def generate_streaming_completion(
        self,
        user_message: str,
    ):
        final_system_prompt = self.default_system_prompt
        max_tokens = self.default_max_tokens

        request_body = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": max_tokens,
            "messages": [
                {"role": "user", "content": user_message}
            ],
            "system": final_system_prompt,
        }

        stream = self.client.invoke_model_with_response_stream(
            modelId=self.model,
            body=json.dumps(request_body),
        )
        for event in stream:
            if "chunk" in event and "bytes" in event["chunk"]:
                chunk_data = json.loads(event["chunk"]["bytes"].decode())
                if "delta" in chunk_data and "text" in chunk_data["delta"]:
                    yield event


if __name__ == "__main__":
    gateway = BedrockLlmGatewayImpl(model="anthropic.claude-3-sonnet-20240229-v1:0")
    result = gateway.generate_completion("Hello, how are you?")
    print(result)