# app/worker.py
import time
from model_loader import ModelLoader
from inference_service import InferenceService
from redis_client import RedisClient
from config import BASE_MODEL, ADAPTER_DIR, HF_TOKEN, REDIS_HOST, REDIS_PORT, TEXT_RESULT_CHANNEL, TEXT_TASK_CHANNEL, BATCH_SIZE, WAIT_TIMEOUT

def main():
    model_loader = ModelLoader(BASE_MODEL, ADAPTER_DIR, HF_TOKEN)
    model_loader.load_model()
    model_loader.load_tokenizer()
    inference_service = InferenceService(model_loader)

    redis_client = RedisClient(
        host=REDIS_HOST,
        port=REDIS_PORT,
        task_channel=TEXT_TASK_CHANNEL,
        result_channel=TEXT_RESULT_CHANNEL
    )

    print("Worker запущен, ожидаем задачи...")

    while True:
        tasks = redis_client.get_tasks(BATCH_SIZE, wait_timeout=WAIT_TIMEOUT)
        if not tasks:
            time.sleep(0.5)
            continue

        texts = [task.text for task in tasks]
        responses = inference_service.generate_batch(texts)
        for task, response in zip(tasks, responses):
            result = {
                "chat_id": task.chat_id,
                "user_id": task.user_id,
                "message_id": task.message_id,
                "text": response
            }
            redis_client.publish_result(result)
            print(f"Обработана задача {task.message_id}")
            
if __name__ == "__main__":
    main()