import asyncio
from typing import List
import ollama

from animuz_core.genai.base import BaseLLM
from animuz_core.genai.prompts import SUMMARIZE_PROMPT, COMBINE_SUMMARIES_PROMPT

class OLlamaClient(BaseLLM):
    def __init__(self, system_prompt: str = None) -> None:
        self.system_prompt = system_prompt
        self.client = ollama.AsyncClient(host="http://127.0.0.1:12084")
        self.message_history = []
        self.temperature = 0.0

    async def get_reply(self, text: str):
        """
        Appends the user's input text to the message history, creates a message using the client with specified parameters, 
        and appends the assistant's response to the message history. Returns the assistant's response text.
        
        Parameters:
            text (str): The user's input text.
        
        Returns:
            str: The assistant's response text.
        """
        self.message_history.append({"role": "user", "content": text})

        response = await self.client.chat(model='llama3', messages=self.message_history)

        self.message_history.append({"role": "assistant", "content": response['message']['content']})
        return response['message']['content']
    
    async def summarize(self, text:str):

        messages = [
            {"role": "system", "content" : SUMMARIZE_PROMPT},
            {"role": "user", "content": text}
                    ]
        response = await self.client.chat(model='llama3:70b', messages=messages)
        print(response)
        print()
        return response['message']['content']
    
    async def summarize_chunks(self, chunks: List[str]):
        res = await asyncio.gather(*[self.summarize(chunk) for chunk in chunks])

        messages = [
            {"role": "system", "content" : COMBINE_SUMMARIES_PROMPT},
            {"role": "user", "content": "\n".join(res)}
                    ]
        response = await self.client.chat(model='llama3:70b', messages=messages)

        return response['message']['content']
