import asyncio
from concurrent.futures import ThreadPoolExecutor
from dotenv import find_dotenv, load_dotenv
from ollama import generate, AsyncClient
from Craft.module.ml.prompt import generate_system, generate_user, generate_system_eng, generate_user_eng
from Craft.module.ml.util import extract_emoji_and_text

load_dotenv(find_dotenv())

class Engine:
    def __init__(self):
        self.key = None
        self.client = AsyncClient(host="http://100.64.0.49:7020")

    async def _generate(self, first_word: str, second_word: str) -> dict:
        return await self.client.generate(
            model="mergegame:latest",
            system=generate_system(first_word, second_word),
            prompt=generate_user(first_word, second_word),
            keep_alive=60*60*24,
            context=None,
            options={
                "seed": 0,
                "temperature": 0.01,
                "top_p": 1,
                "top_k": 0,
                "max_tokens": 16,
                "main_gpu": 0,
            }
        )
        
    async def _generate_eng(self, first_word: str, second_word: str) -> dict:
        system = await generate_system_eng(first_word, second_word)
        prompt = await generate_user_eng(first_word, second_word)
        return generate(
            model="mistral:latest",
            system=system,
            prompt=prompt,
            keep_alive=60*60*24,
            options={
                "seed": 0,
                "temperature": 0.8,
                "top_p": 1,
                "top_k": 0.1,
                "max_tokens": 64,
                "main_gpu": 0,
            }
        )
    
    async def generate(self, first_word: str, second_word: str) -> str:
        #eng_result = await self._generate_eng(first_word, second_word)
        data = await self._generate(first_word, second_word)
        return extract_emoji_and_text(data['response'])