from typing import TYPE_CHECKING, Any, Final, cast

from httpx import Response

from litellm import verbose_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
    _parse_content_for_reasoning,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
    LiteLLMLoggingObj,
)
from litellm.types.llms.bedrock import AmazonDeepSeekR1StreamingResponse
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import (
    ChatCompletionUsageBlock,
    Choices,
    Delta,
    Message,
    ModelResponse,
    ModelResponseStream,
    StreamingChoices,
)

from .amazon_llama_transformation import AmazonLlamaConfig

if TYPE_CHECKING:
    import tiktoken


class AmazonDeepSeekR1Config(AmazonLlamaConfig):
    def transform_response(
        self,
        model: str,
        raw_response: Response,
        model_response: ModelResponse,
        logging_obj: LiteLLMLoggingObj,
        request_data: dict,
        messages: list[AllMessageValues],
        optional_params: dict,
        litellm_params: dict,
        encoding: "tiktoken.Encoding | None",
        api_key: str | None = None,
        json_mode: bool | None = None,
    ) -> ModelResponse:
        """
        Extract the reasoning content, and return it as a separate field in the response.
        """
        response: Final = super().transform_response(
            model,
            raw_response,
            model_response,
            logging_obj,
            request_data,
            messages,
            optional_params,
            litellm_params,
            encoding,
            api_key,
            json_mode,
        )
        prompt: Final = cast(str | None, request_data.get("prompt"))
        message_content: Final = cast(str | None, cast(Choices, response.choices[0]).message.get("content"))
        if prompt and prompt.strip().endswith("<think>") and message_content:
            message_content_with_reasoning_token: Final = "<think>" + message_content
            reasoning, content = _parse_content_for_reasoning(message_content_with_reasoning_token)
            provider_specific_fields: Final = cast(Choices, response.choices[0]).message.provider_specific_fields or {}
            if reasoning:
                provider_specific_fields["reasoning_content"] = reasoning

            message: Final = Message(
                **{
                    **cast(Choices, response.choices[0]).message.model_dump(),
                    "content": content,
                    "provider_specific_fields": provider_specific_fields,
                }
            )
            cast(Choices, response.choices[0]).message = message
        return response


class AmazonDeepseekR1ResponseIterator(BaseModelResponseIterator):
    def __init__(self, streaming_response: Any, sync_stream: bool) -> None:
        super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
        self.has_finished_thinking = False

    def chunk_parser(self, chunk: dict) -> ModelResponseStream:
        """
        Deepseek r1 starts by thinking, then it generates the response.
        """
        try:
            typed_chunk: Final = AmazonDeepSeekR1StreamingResponse(**chunk)
            generated_content = typed_chunk["generation"]
            if generated_content == "</think>" and not self.has_finished_thinking:
                verbose_logger.debug("Deepseek r1: </think> received, setting has_finished_thinking to True")
                generated_content = ""
                self.has_finished_thinking = True

            prompt_token_count: Final = typed_chunk.get("prompt_token_count") or 0
            generation_token_count: Final = typed_chunk.get("generation_token_count") or 0
            usage: Final = ChatCompletionUsageBlock(
                prompt_tokens=prompt_token_count,
                completion_tokens=generation_token_count,
                total_tokens=prompt_token_count + generation_token_count,
            )

            return ModelResponseStream(
                choices=[
                    StreamingChoices(
                        finish_reason=typed_chunk["stop_reason"],
                        delta=Delta(
                            content=(generated_content if self.has_finished_thinking else None),
                            reasoning_content=(generated_content if not self.has_finished_thinking else None),
                        ),
                    )
                ],
                usage=usage,
            )

        except Exception as e:
            raise e
