import copy
import json
import time
from typing import TYPE_CHECKING, Any, Final, cast, get_args

import httpx
from pydantic import TypeAdapter, ValidationError

import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.prompt_templates.factory import (
    cohere_message_pt,
    custom_prompt,
    deepseek_r1_pt,
    prompt_factory,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.request_metadata import (
    bedrock_request_metadata_headers,
    merge_bedrock_invoke_headers,
)
from litellm.llms.custom_httpx.http_handler import (
    AsyncHTTPHandler,
    HTTPHandler,
    _get_httpx_client,
)
from litellm.types.llms.bedrock import GuardrailConfigBlock
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse, Usage
from litellm.utils import CustomStreamWrapper

if TYPE_CHECKING:
    import tiktoken

    from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj

    LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
    LiteLLMLoggingObj = Any

from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM

_GUARDRAIL_CONFIG_VALIDATOR: Final["TypeAdapter[GuardrailConfigBlock]"] = TypeAdapter(GuardrailConfigBlock)

_GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = (
    "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}"
)


def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]":
    try:
        guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config)
    except ValidationError as e:
        raise BedrockError(
            status_code=400,
            message=f"Invalid guardrailConfig={raw_guardrail_config}. Expected format: {_GUARDRAIL_CONFIG_EXPECTED_FORMAT}. Error: {e}",
        )
    if "guardrailIdentifier" not in guardrail_config:
        raise BedrockError(
            status_code=400,
            message=f"guardrailConfig={raw_guardrail_config} is missing 'guardrailIdentifier'. Expected format: {_GUARDRAIL_CONFIG_EXPECTED_FORMAT}",
        )
    trace: Final = guardrail_config.get("trace")
    candidate_headers: Final = {
        "X-Amzn-Bedrock-GuardrailIdentifier": guardrail_config.get("guardrailIdentifier"),
        "X-Amzn-Bedrock-GuardrailVersion": guardrail_config.get("guardrailVersion"),
        "X-Amzn-Bedrock-Trace": trace.upper() if trace is not None else None,
    }
    return {name: value for name, value in candidate_headers.items() if value is not None}


class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
    def __init__(self, **kwargs):
        BaseConfig.__init__(self, **kwargs)
        BaseAWSLLM.__init__(self, **kwargs)

    def get_supported_openai_params(self, model: str) -> list[str]:
        """
        This is a base invoke model mapping. For Invoke - define a bedrock provider specific config that extends this class.
        """
        return [
            "max_tokens",
            "max_completion_tokens",
            "stream",
        ]

    def map_openai_params(
        self,
        non_default_params: dict,
        optional_params: dict,
        model: str,
        drop_params: bool,
    ) -> dict:
        """
        This is a base invoke model mapping. For Invoke - define a bedrock provider specific config that extends this class.
        """
        for param, value in non_default_params.items():
            if param == "max_tokens" or param == "max_completion_tokens":
                optional_params["max_tokens"] = value
            if param == "stream":
                optional_params["stream"] = value
        return optional_params

    def get_complete_url(
        self,
        api_base: str | None,
        api_key: str | None,
        model: str,
        optional_params: dict,
        litellm_params: dict,
        stream: bool | None = None,
    ) -> str:
        """
        Get the complete url for the request
        """
        provider: Final = self.get_bedrock_invoke_provider(model)
        modelId: Final = self.get_bedrock_model_id(
            model=model,
            provider=provider,
            optional_params=optional_params,
        )
        ### SET RUNTIME ENDPOINT ###
        aws_bedrock_runtime_endpoint: Final = optional_params.get(
            "aws_bedrock_runtime_endpoint", None
        )  # https://bedrock-runtime.{region_name}.amazonaws.com
        endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint(
            api_base=api_base,
            aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
            aws_region_name=self._get_aws_region_name(optional_params=optional_params, model=model),
        )

        if (stream is not None and stream is True) and provider != "ai21":
            endpoint_url = f"{endpoint_url}/model/{modelId}/invoke-with-response-stream"
            proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream"
        else:
            endpoint_url = f"{endpoint_url}/model/{modelId}/invoke"
            proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke"

        return endpoint_url

    def sign_request(
        self,
        headers: dict,
        optional_params: dict,
        request_data: dict,
        api_base: str,
        api_key: str | None = None,
        model: str | None = None,
        stream: bool | None = None,
        fake_stream: bool | None = None,
    ) -> tuple[dict, bytes | None]:
        return self._sign_request(
            service_name="bedrock",
            headers=headers,
            optional_params=optional_params,
            request_data=request_data,
            api_base=api_base,
            api_key=api_key,
            model=model,
            stream=stream,
            fake_stream=fake_stream,
        )

    def _apply_config_to_params(self, config: dict, inference_params: dict) -> None:
        """Apply config values to inference_params if not already set."""
        for k, v in config.items():
            if k not in inference_params:
                inference_params[k] = v

    def transform_request(
        self,
        model: str,
        messages: list[AllMessageValues],
        optional_params: dict,
        litellm_params: dict,
        headers: dict,
    ) -> dict:
        ## SETUP ##
        stream: Final = optional_params.pop("stream", None)
        optional_params.pop("stream_chunk_size", None)
        custom_prompt_dict: Final[dict] = litellm_params.pop("custom_prompt_dict", None) or {}
        hf_model_name: Final = litellm_params.get("hf_model_name", None)

        provider: Final = self.get_bedrock_invoke_provider(model)

        prompt, chat_history = self.convert_messages_to_prompt(
            model=hf_model_name or model,
            messages=messages,
            provider=provider,
            custom_prompt_dict=custom_prompt_dict,
        )
        inference_params = copy.deepcopy(optional_params)
        inference_params = {k: v for k, v in inference_params.items() if k not in self.aws_authentication_params}
        request_data: dict = {}
        if provider == "cohere":
            if model.startswith("cohere.command-r"):
                ## LOAD CONFIG
                config = litellm.AmazonCohereChatConfig().get_config()
                self._apply_config_to_params(config, inference_params)
                _data: Final = {"message": prompt, **inference_params}
                if chat_history is not None:
                    _data["chat_history"] = chat_history
                request_data = _data
            else:
                ## LOAD CONFIG
                config = litellm.AmazonCohereConfig.get_config()
                self._apply_config_to_params(config, inference_params)
                if stream is True:
                    inference_params["stream"] = True  # cohere requires stream = True in inference params
                request_data = {"prompt": prompt, **inference_params}
        elif provider == "anthropic":
            transformed_request: Final = litellm.AmazonAnthropicClaudeConfig().transform_request(
                model=model,
                messages=messages,
                optional_params=optional_params,
                litellm_params=litellm_params,
                headers=headers,
            )

            return transformed_request
        elif provider == "nova":
            return litellm.AmazonInvokeNovaConfig().transform_request(
                model=model,
                messages=messages,
                optional_params=optional_params,
                litellm_params=litellm_params,
                headers=headers,
            )
        elif provider == "ai21":
            ## LOAD CONFIG
            config = litellm.AmazonAI21Config.get_config()
            self._apply_config_to_params(config, inference_params)
            request_data = {"prompt": prompt, **inference_params}
        elif provider == "mistral":
            ## LOAD CONFIG
            config = litellm.AmazonMistralConfig.get_config()
            self._apply_config_to_params(config, inference_params)
            request_data = {"prompt": prompt, **inference_params}
        elif provider == "amazon":  # amazon titan
            ## LOAD CONFIG
            config = litellm.AmazonTitanConfig.get_config()
            self._apply_config_to_params(config, inference_params)
            request_data = {
                "inputText": prompt,
                "textGenerationConfig": inference_params,
            }
        elif provider == "meta" or provider == "llama" or provider == "deepseek_r1":
            ## LOAD CONFIG
            config = litellm.AmazonLlamaConfig.get_config()
            self._apply_config_to_params(config, inference_params)
            request_data = {"prompt": prompt, **inference_params}
        elif provider == "twelvelabs":
            return litellm.AmazonTwelveLabsPegasusConfig().transform_request(
                model=model,
                messages=messages,
                optional_params=optional_params,
                litellm_params=litellm_params,
                headers=headers,
            )
        elif provider == "openai":
            # OpenAI imported models use OpenAI Chat Completions format
            return litellm.AmazonBedrockOpenAIConfig().transform_request(
                model=model,
                messages=messages,
                optional_params=optional_params,
                litellm_params=litellm_params,
                headers=headers,
            )
        else:
            raise BedrockError(
                status_code=404,
                message=f"Bedrock Invoke HTTPX: Unknown provider={provider}, model={model}. Try calling via converse route - `bedrock/converse/<model>`.",
            )

        return request_data

    def transform_response(
        self,
        model: str,
        raw_response: httpx.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:
        try:
            completion_response: Final = raw_response.json()
        except Exception:
            raise BedrockError(message=raw_response.text, status_code=raw_response.status_code)
        verbose_logger.debug(
            "bedrock invoke response % s",
            json.dumps(completion_response, indent=4, default=str),
        )
        provider: Final = self.get_bedrock_invoke_provider(model)
        outputText: str | None = None
        try:
            if provider == "cohere":
                if "text" in completion_response:
                    outputText = completion_response["text"]
                elif "generations" in completion_response:
                    outputText = completion_response["generations"][0]["text"]
                    model_response.choices[0].finish_reason = map_finish_reason(
                        completion_response["generations"][0]["finish_reason"]
                    )
            elif provider == "anthropic":
                return litellm.AmazonAnthropicClaudeConfig().transform_response(
                    model=model,
                    raw_response=raw_response,
                    model_response=model_response,
                    logging_obj=logging_obj,
                    request_data=request_data,
                    messages=messages,
                    optional_params=optional_params,
                    litellm_params=litellm_params,
                    encoding=encoding,
                    api_key=api_key,
                    json_mode=json_mode,
                )
            elif provider == "nova":
                return litellm.AmazonInvokeNovaConfig().transform_response(
                    model=model,
                    raw_response=raw_response,
                    model_response=model_response,
                    logging_obj=logging_obj,
                    request_data=request_data,
                    messages=messages,
                    optional_params=optional_params,
                    litellm_params=litellm_params,
                    encoding=encoding,
                )
            elif provider == "twelvelabs":
                return litellm.AmazonTwelveLabsPegasusConfig().transform_response(
                    model=model,
                    raw_response=raw_response,
                    model_response=model_response,
                    logging_obj=logging_obj,
                    request_data=request_data,
                    messages=messages,
                    optional_params=optional_params,
                    litellm_params=litellm_params,
                    encoding=encoding,
                    api_key=api_key,
                    json_mode=json_mode,
                )
            elif provider == "ai21":
                outputText = completion_response.get("completions")[0].get("data").get("text")
            elif provider == "meta" or provider == "llama" or provider == "deepseek_r1":
                outputText = completion_response["generation"]
            elif provider == "mistral":
                outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response)
            else:  # amazon titan
                outputText = completion_response.get("results")[0].get("outputText")
        except Exception as e:
            raise BedrockError(
                message=f"Error processing={raw_response.text}, Received error={e}",
                status_code=422,
            )

        try:
            if (
                outputText is not None
                and len(outputText) > 0
                and hasattr(model_response.choices[0], "message")
                and getattr(model_response.choices[0].message, "tool_calls", None) is None
            ):
                model_response.choices[0].message.content = outputText
            elif (
                hasattr(model_response.choices[0], "message")
                and getattr(model_response.choices[0].message, "tool_calls", None) is not None
            ):
                pass
            else:
                raise Exception()
        except Exception as e:
            raise BedrockError(
                message=f"Error parsing received text={outputText}.\nError-{e}",
                status_code=raw_response.status_code,
            )

        ## CALCULATING USAGE - bedrock returns usage in the headers
        bedrock_input_tokens: Final = raw_response.headers.get("x-amzn-bedrock-input-token-count", None)
        bedrock_output_tokens: Final = raw_response.headers.get("x-amzn-bedrock-output-token-count", None)

        prompt_tokens: Final = int(bedrock_input_tokens or litellm.token_counter(messages=messages))

        completion_tokens: Final = int(
            bedrock_output_tokens
            or litellm.token_counter(
                text=model_response.choices[0].message.content,
                count_response_tokens=True,
            )
        )

        model_response.created = int(time.time())
        model_response.model = model
        usage: Final = Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )
        setattr(model_response, "usage", usage)

        return model_response

    def validate_environment(
        self,
        headers: dict,
        model: str,
        messages: list[AllMessageValues],
        optional_params: dict,
        litellm_params: dict,
        api_key: str | None = None,
        api_base: str | None = None,
    ) -> dict:
        raw_guardrail_config: Final = optional_params.pop("guardrailConfig", None)
        guardrail_headers: Final = (
            ()
            if raw_guardrail_config is None
            else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items())
        )
        owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
        return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)

    def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
        return BedrockError(status_code=status_code, message=error_message)

    @track_llm_api_timing()
    async def get_async_custom_stream_wrapper(
        self,
        model: str,
        custom_llm_provider: str,
        logging_obj: LiteLLMLoggingObj,
        api_base: str,
        headers: dict,
        data: dict,
        messages: list,
        client: AsyncHTTPHandler | None = None,
        json_mode: bool | None = None,
        signed_json_body: bytes | None = None,
    ) -> CustomStreamWrapper:
        completion_stream, response_headers = await make_call(
            client=client,
            api_base=api_base,
            headers=headers,
            data=json.dumps(data),
            model=model,
            messages=messages,
            logging_obj=logging_obj,
            fake_stream=True if "ai21" in api_base else False,
            bedrock_invoke_provider=self.get_bedrock_invoke_provider(model),
            json_mode=json_mode,
        )
        streaming_response: Final = CustomStreamWrapper(
            completion_stream=completion_stream,
            model=model,
            custom_llm_provider="bedrock",
            logging_obj=logging_obj,
            _response_headers=response_headers,
        )
        return streaming_response

    @track_llm_api_timing()
    def get_sync_custom_stream_wrapper(
        self,
        model: str,
        custom_llm_provider: str,
        logging_obj: LiteLLMLoggingObj,
        api_base: str,
        headers: dict,
        data: dict,
        messages: list,
        client: HTTPHandler | AsyncHTTPHandler | None = None,
        json_mode: bool | None = None,
        signed_json_body: bytes | None = None,
    ) -> CustomStreamWrapper:
        sync_client: Final = (
            _get_httpx_client(params={}) if client is None or isinstance(client, AsyncHTTPHandler) else client
        )
        completion_stream, response_headers = make_sync_call(
            client=sync_client,
            api_base=api_base,
            headers=headers,
            data=json.dumps(data),
            signed_json_body=signed_json_body,
            model=model,
            messages=messages,
            logging_obj=logging_obj,
            fake_stream=True if "ai21" in api_base else False,
            bedrock_invoke_provider=self.get_bedrock_invoke_provider(model),
            json_mode=json_mode,
        )
        streaming_response: Final = CustomStreamWrapper(
            completion_stream=completion_stream,
            model=model,
            custom_llm_provider="bedrock",
            logging_obj=logging_obj,
            _response_headers=response_headers,
        )
        return streaming_response

    @property
    def has_custom_stream_wrapper(self) -> bool:
        return True

    @property
    def supports_stream_param_in_request_body(self) -> bool:
        """
        Bedrock invoke does not allow passing `stream` in the request body.
        """
        return False

    @staticmethod
    def get_bedrock_invoke_provider(
        model: str,
    ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None:
        """
        Helper function to get the bedrock provider from the model

        handles 4 scenarios:
        1. model=invoke/anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic`
        2. model=anthropic.claude-3-5-sonnet-20240620-v1:0 -> Returns `anthropic`
        3. model=llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n -> Returns `llama`
        4. model=us.amazon.nova-pro-v1:0 -> Returns `nova`
        """
        if model.startswith("invoke/"):
            model = model.replace("invoke/", "", 1)

        # Special case: Check for "nova" in model name first (before "amazon")
        # This handles amazon.nova-* models which would otherwise match "amazon" (Titan)
        if "nova" in model.lower():
            if "nova" in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
                return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, "nova")

        _split_model: Final = model.split(".")[0]
        if _split_model in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
            return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, _split_model)

        # If not a known provider, check for pattern with two slashes
        provider = AmazonInvokeConfig._get_provider_from_model_path(model)
        if provider is not None:
            return provider

        for provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
            if provider in model:
                return provider
        return None

    @staticmethod
    def _get_provider_from_model_path(
        model_path: str,
    ) -> litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL | None:
        """
        Helper function to get the provider from a model path with format: provider/model-name

        Args:
            model_path (str): The model path (e.g., 'llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n' or 'anthropic/model-name')

        Returns:
            Optional[str]: The provider name, or None if no valid provider found
        """
        parts: Final = model_path.split("/")
        if len(parts) >= 1:
            provider: Final = parts[0]
            if provider in get_args(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL):
                return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
        return None

    def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> tuple[str, list | None]:
        # handle anthropic prompts and amazon titan prompts
        prompt = ""
        chat_history: list | None = None
        ## CUSTOM PROMPT
        if model in custom_prompt_dict:
            # check if the model has a registered custom prompt
            model_prompt_details: Final = custom_prompt_dict[model]
            prompt = custom_prompt(
                role_dict=model_prompt_details["roles"],
                initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
                final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
                messages=messages,
            )
            return prompt, None
        ## ELSE
        if (
            provider == "anthropic"
            or provider == "amazon"
            or provider == "mistral"
            or provider == "meta"
            or provider == "llama"
        ):
            prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="bedrock")
        elif provider == "cohere":
            prompt, chat_history = cohere_message_pt(messages=messages)
        elif provider == "deepseek_r1":
            prompt = deepseek_r1_pt(messages=messages)
        else:
            prompt = ""
            for message in messages:
                if "role" in message:
                    if message["role"] == "user":
                        prompt += f"{message['content']}"
                    else:
                        prompt += f"{message['content']}"
                else:
                    prompt += f"{message['content']}"
        return prompt, chat_history
