from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate
from langchain_groq import ChatGroq
from langchain_huggingface import ChatHuggingFace
from langchain_huggingface import HuggingFaceEndpoint
from langchain_google_genai import ChatGoogleGenerativeAI
from dotenv import load_dotenv
from huggingface_hub import login
import os

# Load environment variables and authenticate
load_dotenv()
login(token=os.environ.get("HUGGING_FACE_API_KEY", ""))
os.environ['CURL_CA_BUNDLE'] = ''

class Bot:
    def __init__(self):
        # Updated model lists: Remove gated or unsupported IDs        
        self.groq_models = [
    'gemma2-9b-it',
    'llama-3.3-70b-versatile',
    'llama-3.1-8b-instant',
    'meta-llama/llama-guard-4-12b'
        ]
        self.hf_models = [
            "01-ai/Yi-1.5-34B-Chat",
            #"google/gemma-1.1-2b-it"
        ]
        # Use supported Google GenAI model names
        self.google_models = [
            "gemini-pro",
            "gemini-pro-vision"
        ]
        # Master list for sampling (only include accessible models)
        self.models = self.google_models + self.hf_models + self.groq_models

    def call_groq(self, model, temp=0.7, given_prompt="Hi"):
        try:
            llm = ChatGroq(model=model, temperature=temp)
            prompt = ChatPromptTemplate.from_messages([
                ("system", "You are a helpful assistant."),
                ("human", "{text}")
            ])
            chain = prompt | llm | StrOutputParser()
            return chain.invoke({"text": given_prompt})
        except Exception as e:
            return f"⚠️ [Groq:{model}] {str(e)}"

    def call_hf(self, model, temp=0.7, given_prompt="Hi"):
        try:
            llm = HuggingFaceEndpoint(repo_id=model, temperature=temp)
            chat = ChatHuggingFace(llm=llm, verbose=True)
            template = """
You are a helpful assistant.
User: {query}
Answer:
"""
            prompt = PromptTemplate(template=template, input_variables=["query"])
            chain = prompt | chat | StrOutputParser()
            return chain.invoke({"query": given_prompt})
        except Exception as e:
            return f"⚠️ [HF:{model}] {str(e)}"

    def call_google(self, model, temp=0.7, given_prompt="Hi"):
        try:
            gm = ChatGoogleGenerativeAI(model=model, temperature=temp)
            prompt = ChatPromptTemplate.from_messages([("system", "You are a helpful assistant."), ("human", "{text}")])
            chain = prompt | gm | StrOutputParser()
            return chain.invoke({"text": given_prompt})
        except Exception as e:
            return f"⚠️ [Google:{model}] {str(e)}"

    def response(self, model, prompt="Hi", temperature=0.7):
        # Route to the correct provider and catch errors
        try:
            if model in self.groq_models:
                return self.call_groq(model, temp=temperature, given_prompt=prompt)
            if model in self.hf_models:
                return self.call_hf(model, temp=temperature, given_prompt=prompt)
            if model in self.google_models:
                return self.call_google(model, temp=temperature, given_prompt=prompt)
            return f"❌ Unsupported model: {model}"
        except Exception as e:
            return f"⚠️ Skipping {model} due to error: {str(e)}"