import os
import time
import logging
from pathlib import Path
from typing import Dict, Any, Type, Union, Optional
import json
from pydantic import BaseModel
from openai import OpenAI, OpenAIError
from litellm import completion
import litellm


# litellm._turn_on_debug()


class LlmConfig(BaseModel):
    model: str = "gpt-o3-mini"  # "gpt-4o-2024-08-06"
    max_retries: int = 30
    temperature: float = 0.0


class LlmClient:
    """Only support OpenAI now"""

    def __init__(
        self,
        system_prompt_template: Union[str, Path],
        llm_config: Optional[LlmConfig] = None,
    ):
        self.config = llm_config or LlmConfig()
        self.client = OpenAI(max_retries=self.config.max_retries)
        self.model = self.config.model
        self.temperature = self.config.temperature
        self.system_prompt = self._load_file(system_prompt_template)

    def _load_file(self, file_path: Union[str, Path]) -> str:
        try:
            path = Path(file_path)
            with open(path, "r") as f:
                return f.read()
        except FileNotFoundError:
            raise RuntimeError(f"File not found: {str(file_path)}")
        except IOError:
            raise RuntimeError(f"Failed to read file: {str(file_path)}")

    def update_system_prompt(self, new_system_prompt: str):
        self.system_prompt = new_system_prompt

    def deprecated_generate(
        self, user_prompt: Dict[str, Any], response_format: Type[BaseModel] = None
    ) -> Union[str, BaseModel]:
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": json.dumps(user_prompt)},
        ]
        try:
            if response_format is None:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=self.temperature,
                )
                logging.info(f"response content: {response.choices[0].message.content}")
                return response.choices[0].message.content
            else:
                response = self.client.beta.chat.completions.parse(
                    model=self.model,
                    messages=messages,
                    response_format=response_format,
                    temperature=self.temperature,
                )
                logging.info(f"response parsed: {response.choices[0].message.parsed}")
                return response.choices[0].message.parsed
        except OpenAIError as e:
            raise RuntimeError(f"OpenAI API error: {e}\nRequest messages: {messages}")
        except Exception as e:
            raise RuntimeError(
                f"Failed to generate response: {e}\nRequest messages: {messages}"
            )

    def generate(
        self, user_prompt: Dict[str, Any], response_format: Type[BaseModel] = None
    ) -> Union[str, BaseModel]:
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": json.dumps(user_prompt)},
        ]
        time.sleep(5)
        self.temperature = 1.0 if "o3-mini" in self.model else 0.0
        try:
            if response_format is None:
                if self.model.startswith("bedrock/"):
                    # there are 2 ways to use bedrock
                    # 1. without inference profile id: set model like `bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0`
                    # 2. with inference profile id: 
                    #   2.1 set model like `bedrock/converse/arn:aws:bedrock:us-east-1:841162690310:application-inference-profile/a0a0a0a0a0a0`
                    #   2.2 set model_id like `arn:aws:bedrock:us-east-1:841162690310:application-inference-profile/a0a0a0a0a0a0` and set the parent model like `bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0`
                    response = completion(
                        model=self.model,
                        messages=messages,
                        temperature=self.temperature,
                        num_retries=self.config.max_retries,
                        aws_access_key_id=os.getenv("AWS_AK"),
                        aws_secret_access_key=os.getenv("AWS_SK"),
                        aws_session_token=None,
                        aws_role_name=None,
                        aws_session_name=None,
                        aws_profile_name=None,
                        aws_sts_endpoint=None,
                        aws_region_name="us-east-1",
                        aws_web_identity_token=None,
                    )
                else:
                    response = completion(
                        model=self.model,
                        messages=messages,
                        temperature=self.temperature,
                        num_retries=self.config.max_retries,
                    )
                logging.info(f"response content: {response.choices[0].message.content}")
                return response.choices[0].message.content
            else:
                if self.model.startswith("bedrock/"):
                    # there are 2 ways to use bedrock
                    # 1. without inference profile id: set model like `bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0`
                    # 2. with inference profile id: 
                    #   2.1 set model like `bedrock/converse/arn:aws:bedrock:us-east-1:841162690310:application-inference-profile/a0a0a0a0a0a0`
                    #   2.2 set model_id like `arn:aws:bedrock:us-east-1:841162690310:application-inference-profile/a0a0a0a0a0a0` and set the parent model like `bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0`
                    response = completion(
                        model=self.model,
                        messages=messages,
                        response_format=response_format,
                        temperature=self.temperature,
                        num_retries=self.config.max_retries,
                        aws_access_key_id=os.getenv("AWS_AK"),
                        aws_secret_access_key=os.getenv("AWS_SK"),
                        aws_session_token=None,
                        aws_role_name=None,
                        aws_session_name=None,
                        aws_profile_name=None,
                        aws_sts_endpoint=None,
                        aws_region_name="us-east-1",
                        aws_web_identity_token=None,
                    )
                else:
                    response = completion(
                        model=self.model,
                        messages=messages,
                        response_format=response_format,
                        temperature=self.temperature,
                        num_retries=self.config.max_retries,
                    )
                logging.info(f"response parsed: {json.loads(response.choices[0].message.content)}")
                return response_format(**json.loads(response.choices[0].message.content))
        except OpenAIError as e:
            raise RuntimeError(f"OpenAI API error: {e}\nRequest messages: {messages}")
        except Exception as e:
            raise RuntimeError(
                f"Failed to generate response: {e}\nRequest messages: {messages}"
            )