from typing import List, Dict, Any
import json
import asyncio
import aiohttp
import time
import re



model_url = "http://192.168.130.206:4001/v1"
model = "google/gemma-3-27b-it"


class LLMModel:
    def __init__(self):

        self.instruction = """
            You are a helpful assistant that help to me to modify and change the input question.
            I will give you a question and its text and you must replace the words of question with synonyms or similar words.

            ## Important:
            - replace the words of question with synonyms or similar words.

            -the question must be in persian language.

            return the question nothing else.
            """

    async def run_llm(self, session, question, text):
        """
        Run the llm model.
        Args:
            session: The session to use for the request.
            question: The question to evaluate the text.
            text: The text to evaluate.
        Returns:
            The result of the text.
        """
        headers = {"Content-Type": "application/json"}

        input_message = f"""{{"question": "{question}", "text": "{text}"}}"""
        messages = [{"role": "system", "content": self.instruction}, {"role": "user", "content": input_message}]

        payload = {
            "model": model,
            "messages": messages,
            "max_tokens": 100
        }
        try:
            async with session.post(model_url + "/chat/completions", headers=headers, json=payload) as resp:
                resp.raise_for_status()
                response = await resp.json()
            
            result = response['choices'][0]['message']['content']
            print(f"question: {question}")
            print(f"result: {result}")
            print("--------------------------------")
            
            return result

        except Exception as e:
            try:
                print(f"Error in llm model {response}: {e}")
            except:
                print(f"Error in llm model: {e}")
            return ""


    async def run_llm_async(self, question_list, text_list):
        """
        Send all chunk requests concurrently.
        Args:
            question_list: The list of questions.
            text_list: The list of texts.
        Returns:
            The list of results.
        """
        async with aiohttp.ClientSession() as session:
            tasks = [self.run_llm(session, question, text) for question, text in zip(question_list, text_list)]
            results = await asyncio.gather(*tasks)
        return results


    def modify_question_llm(self, query_list: List[str], text_list: List[str]) -> List[Dict[str, Any]]:
        """
        Modify question of the documents based on the query using the LLM model.
        Args:
            query_list: The list of queries.
            text_list: The list of texts.
        Returns:
            The list of modified questions.
        """
        if not text_list:
            return []

        start_time = time.time()
        results = asyncio.run(self.run_llm_async(query_list, text_list))
        end_time = time.time()
        # print(f"Time taken for llm model: {end_time - start_time} seconds")
        
        return results