from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain.prompts import PromptTemplate
from langchain.chains.question_answering import load_qa_chain
from langchain_google_genai import ChatGoogleGenerativeAI
import google.generativeai as genai
import os
from dotenv import load_dotenv


def list_available_models():
    models = genai.models.list_models()

    print("Available models:")
    for model in models:
        print(f"Name: {model.name}")
        print(f"Description: {model.description}")
        print(f"Supported methods: {', '.join(model.supported_methods)}")
        print("\n")

def get_response(file, query):
    # Load environment variables
    load_dotenv()

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=400)
    context = '\n\n'.join(str(p.page_content) for p in file)
    data = text_splitter.split_text(context)

    # Specify the correct model name based on your requirements
    model_name = 'models/chat-bison-001'
    
    # Specify the API key directly in the code
    google_api_key = os.getenv("GOOGLE_API_KEY")

    embeddings = GoogleGenerativeAIEmbeddings(model=model_name, google_api_key=google_api_key)

    searcher = Chroma.from_texts(data, embeddings).as_retriever()

    ques = 'Which country has maximum GDP?'
    records = searcher.get_relevent_documents(ques)

    prompt_template = """
        You have to give the correct answer to the question from the provided context and make sure you give all details\n
        Context: {context}\n
        Question: {question}\n

        Answer:
    """
    prompt = PromptTemplate(template=prompt_template, input_variable=['context', 'question'])

    model = ChatGoogleGenerativeAI(model=model_name, temperature=0.5)

    chain = load_qa_chain(model, chain_type='stuff', prompt=prompt)

    response = chain(
        {
            'input_document': records,
            'question': query
        },
        return_only_output=True
    )

    return response['output_text']