"""Amazon Bedrock LLM implementation using boto3."""

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator

from roohai.base import LLMModel


class BedrockLLM(LLMModel):
    META = {
        "display_name": "Amazon Bedrock (Claude)",
        "description": "Anthropic Claude via Amazon Bedrock. Cloud-based, uses AWS credentials.",
        "kind": "cloud",
        "pip_dependencies": ["boto3"],
        "hf_models": [],
        "estimated_download_mb": 0,
        "config_fields": [
            {
                "name": "auth_mode",
                "label": "Auth Mode",
                "type": "select",
                "required": False,
                "default": "environment",
                "options": ["environment", "access_keys"],
            },
            {
                "name": "aws_access_key_id",
                "label": "AWS Access Key ID",
                "type": "secret",
                "required": False,
                "env_var": "AWS_ACCESS_KEY_ID",
                "depends_on": {"auth_mode": "access_keys"},
            },
            {
                "name": "aws_secret_access_key",
                "label": "AWS Secret Access Key",
                "type": "secret",
                "required": False,
                "env_var": "AWS_SECRET_ACCESS_KEY",
                "depends_on": {"auth_mode": "access_keys"},
            },
            {
                "name": "region",
                "label": "AWS Region",
                "type": "text",
                "required": False,
                "default": "us-east-1",
            },
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",
                "required": False,
                "default": "anthropic.claude-3-haiku-20240307-v1:0",
            },
        ],
    }

    def __init__(
        self,
        model_id: str = "anthropic.claude-3-haiku-20240307-v1:0",
        region: str = "us-east-1",
        auth_mode: str = "environment",
        aws_access_key_id: str | None = None,
        aws_secret_access_key: str | None = None,
    ) -> None:
        self.model_id = model_id
        self.region = region
        self.auth_mode = auth_mode
        self.aws_access_key_id = aws_access_key_id
        self.aws_secret_access_key = aws_secret_access_key
        self._client = None

    def load(self) -> None:
        self._ensure_client()

    def unload(self) -> None:
        if self._client is not None:
            self._client.close()
            self._client = None

    @property
    def is_loaded(self) -> bool:
        return self._client is not None

    def _ensure_client(self) -> None:
        if self._client is None:
            import boto3

            if self.auth_mode == "access_keys" and self.aws_access_key_id and self.aws_secret_access_key:
                self._client = boto3.client(
                    "bedrock-runtime",
                    region_name=self.region,
                    aws_access_key_id=self.aws_access_key_id,
                    aws_secret_access_key=self.aws_secret_access_key,
                )
            else:
                self._client = boto3.client("bedrock-runtime", region_name=self.region)

    def _invoke(self, message: str, history: list[dict] | None, system_prompt: str | None) -> str:
        self._ensure_client()
        messages = []
        if history:
            messages.extend(history)
        messages.append({"role": "user", "content": message})

        body: dict = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 1024,
            "messages": messages,
        }
        if system_prompt:
            body["system"] = system_prompt

        response = self._client.invoke_model(
            modelId=self.model_id,
            contentType="application/json",
            accept="application/json",
            body=json.dumps(body),
        )
        result = json.loads(response["body"].read())
        return result["content"][0]["text"]

    def _invoke_stream(self, message: str, history: list[dict] | None, system_prompt: str | None):
        """Invoke Bedrock with streaming and yield text chunks synchronously."""
        self._ensure_client()
        messages = []
        if history:
            messages.extend(history)
        messages.append({"role": "user", "content": message})

        body: dict = {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 1024,
            "messages": messages,
        }
        if system_prompt:
            body["system"] = system_prompt

        response = self._client.invoke_model_with_response_stream(
            modelId=self.model_id,
            contentType="application/json",
            accept="application/json",
            body=json.dumps(body),
        )

        for event in response["body"]:
            chunk_bytes = event.get("chunk", {}).get("bytes")
            if not chunk_bytes:
                continue
            chunk_data = json.loads(chunk_bytes)
            if chunk_data.get("type") == "content_block_delta":
                delta = chunk_data.get("delta", {})
                text = delta.get("text", "")
                if text:
                    yield text

    async def chat(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> str:
        return await asyncio.to_thread(self._invoke, message, history, system_prompt)

    async def chat_stream(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> AsyncIterator[str]:
        """Yield text tokens as they stream from Bedrock."""
        loop = asyncio.get_event_loop()
        queue: asyncio.Queue[str | None] = asyncio.Queue()

        def _run():
            try:
                for chunk in self._invoke_stream(message, history, system_prompt):
                    loop.call_soon_threadsafe(queue.put_nowait, chunk)
            finally:
                loop.call_soon_threadsafe(queue.put_nowait, None)

        fut = loop.run_in_executor(None, _run)
        while True:
            item = await queue.get()
            if item is None:
                break
            yield item
        await fut
