# Standard Import
import os
import sys

# Lib Import
import anthropic
from anthropic import AnthropicVertex
import openai
import google.generativeai as genai
from retry import retry
import vertexai
from vertexai.generative_models import GenerativeModel
import vertexai.preview.generative_models as generative_models

# App Import
from src.constants import (
    LLM,
    Task,
    CLAUDE_LOCATION,
    CLAUDE_MODEL_NAMES,
    OPENAI_MODEL_NAMES,
    GOOGLE_MODEL_NAMES,
    GOOGLE_GENERATION_CONFIG,
    GOOGLE_LOCATION,
    GOOGLE_VERBALIZATION_CONFIG,
    VERTEX_AI_PROJECT_ID
)

# Configuration Block

openai_client = openai.OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY")
)  # OpenAI client initialization

#genai.configure(
#    api_key=os.environ["GEMINI_API_KEY"]
#)  # Google Generative AI client initialization

anthropic_client = anthropic.Anthropic(
    api_key=os.environ["ANTHROPIC_API_KEY"]
)  # Anthropic client initialization


class LLMManager:
    """
    Manager for interfacing with different Language Models.

    This class manages interactions with various Language Models based on the provided model name.
    It provides methods to interact with different Language Models through a unified interface.
    """

    def __init__(self, llm, task, system_prompt, user_prompt):
        """
        Initialize LLMManager instance.

        Args:
            llm (str): LLM name.
            subject (str): Subject name.
        """
        self.llm = llm
        self.system_prompt = system_prompt
        self.user_prompt = user_prompt
        self.task = task
        self.model_dispatcher = {
            LLM.GPT4.value: self.call_openai,
            LLM.Claude.value: self.call_claude,
            LLM.Gemini.value: self.call_gemini,
        }

    def call_llm(self):
        """
        Call the appropriate Language Model based on the provided model name.

        Returns:
            str: Response from the Language Model.
        """
        response = self.model_dispatcher[self.llm]()
        return response

    @retry(tries=2, delay=1)
    def call_openai(self):
        """
        Call OpenAI language model.

        Returns:
            str: Response from the OpenAI language model.
        """
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": self.user_prompt},
        ]

        if self.task == Task.Verbalization.value:
            print(OPENAI_MODEL_NAMES[self.llm])
            response = openai_client.chat.completions.create(
                messages=messages,
                model=OPENAI_MODEL_NAMES[self.llm],
                temperature=0.1
            )
        elif self.task == Task.Generation.value:
            response = openai_client.chat.completions.create(
                messages=messages,
                model=OPENAI_MODEL_NAMES[self.llm],
                temperature=0.1
            )

        return response.choices[0].message.content

    @retry(tries=2, delay=1)
    def call_claude(self):
        """
        Call Claude language model.

        Returns:
            str: Response from the Claude language model.
        """
        client = AnthropicVertex(region=CLAUDE_LOCATION, project_id=VERTEX_AI_PROJECT_ID)
        messages = [{"role": "user", "content": self.user_prompt}]

        response = client.messages.create(
            max_tokens=4096,
            temperature=0.1,
            messages=messages,
            model=CLAUDE_MODEL_NAMES[self.llm]
        )

        return response.content[0].text

    @retry(tries=2, delay=1)
    def call_gemini(self):
        """
        Call Gemini language model.

        Returns:
            str: Response from the Gemini language model.
        """
        model_name = GOOGLE_MODEL_NAMES[self.llm]
        if self.task == Task.Verbalization.value:
            model_config = GOOGLE_VERBALIZATION_CONFIG
        elif self.task == Task.Generation.value:
            model_config = GOOGLE_GENERATION_CONFIG

        vertexai.init(project=VERTEX_AI_PROJECT_ID, location=GOOGLE_LOCATION)

        model = GenerativeModel(model_name)
        safety_settings = {
            generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
            generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
            generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_ONLY_HIGH,
            generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
        }

        response = model.generate_content(
            [self.user_prompt],
            generation_config=model_config,
            safety_settings=safety_settings,
        )

        return response.text