import asyncio
from concurrent.futures import ThreadPoolExecutor
from dotenv import find_dotenv, load_dotenv
from ollama import generate
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

    def _generate(self, first_word: str, second_word: str, eng_result: str) -> dict:
        return generate(
            model="infcraft:latest",
            system=generate_system(first_word, second_word, eng_result),
            prompt=generate_user(first_word, second_word),
            keep_alive=60*60*24,
            context=None,
            options={
                "seed": 0,
                "temperature": 0.4,
                "top_p": 0.85,
                "top_k": 0.1,
                "max_tokens": 32,
                "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:
        loop = asyncio.get_running_loop()
        eng_result = await self._generate_eng(first_word, second_word)
        with ThreadPoolExecutor() as executor:
            data = await loop.run_in_executor(executor, self._generate, first_word, second_word, eng_result['response'])
        return extract_emoji_and_text(data['response'])