import logging
import re
from typing import AsyncGenerator

import httpx
from fastapi import HTTPException

from llm_bridge.client.implementations.claude.claude_token_counter import count_claude_input_tokens, \
    count_claude_output_tokens
from llm_bridge.client.model_client.claude_client import ClaudeClient
from llm_bridge.type.chat_response import ChatResponse
from llm_bridge.type.serializer import serialize


class PrintingStatus:
    Start = "start"
    Thought = "thought"
    Response = "response"


class StreamClaudeClient(ClaudeClient):
    async def generate_stream_response(self) -> AsyncGenerator[ChatResponse, None]:
        try:
            logging.info(f"messages: {self.messages}")

            try:
                async with self.client.beta.messages.stream(
                    model=self.model,
                    max_tokens=self.max_tokens,
                    temperature=self.temperature,
                    system=self.system,
                    messages=serialize(self.messages),
                    thinking=self.thinking,
                    betas=self.betas,
                ) as stream:
                    printing_status = PrintingStatus.Start

                    async for event in stream:
                        text = ""
                        if event.type == "content_block_start":
                            continue
                        elif event.type == "content_block_delta":
                            if event.delta.type == "thinking_delta":
                                if printing_status == PrintingStatus.Start:
                                    text += "# Model Thought:\n\n"
                                    printing_status = PrintingStatus.Thought
                                text += event.delta.thinking
                            elif event.delta.type == "text_delta":
                                if printing_status == PrintingStatus.Thought:
                                    text += "\n\n# Model Response:\n\n"
                                    printing_status = PrintingStatus.Response
                                text += event.delta.text

                        chat_response = ChatResponse(text=text)
                        output_tokens = await count_claude_output_tokens(
                            client=self.client,
                            model=self.model,
                            chat_response=chat_response,
                        )
                        yield ChatResponse(
                            text=text,
                            input_tokens=self.input_tokens,
                            output_tokens=output_tokens,
                        )

            except Exception as e:
                logging.exception(e)
                yield ChatResponse(error=repr(e))

        except httpx.HTTPStatusError as e:
            status_code = e.response.status_code
            text = e.response.text
            raise HTTPException(status_code=status_code, detail=text)
        except Exception as e:
            logging.exception(e)
            match = re.search(r'\d{3}', str(e))
            if match:
                error_code = int(match.group(0))
            else:
                error_code = 500

            raise HTTPException(status_code=error_code, detail=str(e))
