import os
import spaces
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from threading import Thread
from queue import Queue, Empty
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model_id = "meta-llama/Meta-Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ.get("MY_API_LLAMA_3_1"))

model = None
model_load_queue = Queue()

def load_model():
    global model
    try:
        if model is None:
            logger.info("Loading model...")
            model = AutoModelForCausalLM.from_pretrained(
                model_id,
                token=os.environ.get("MY_API_LLAMA_3_1"),
                torch_dtype=torch.bfloat16,
                device_map="auto",
                low_cpu_mem_usage=True,
                load_in_8bit=True
            )
            logger.info("Model loaded successfully")
        model_load_queue.put(model)
    except Exception as e:
        logger.error(f"Error loading model: {str(e)}")
        model_load_queue.put(None)

@spaces.GPU(duration=120)
def generate_response(chat, kwargs):
    global model
    try:
        if model is None:
            logger.info("Starting model loading thread")
            Thread(target=load_model).start()
            model = model_load_queue.get(timeout=120)
            if model is None:
                return "Nie udało się załadować modelu. Proszę spróbować ponownie później."

        logger.info("Preparing input for generation")
        inputs = tokenizer(chat, return_tensors="pt").to(model.device)
        streamer = TextIteratorStreamer(tokenizer, timeout=120., skip_prompt=True, skip_special_tokens=True)

        if 'seed' in kwargs:
            del kwargs['seed']

        generation_kwargs = dict(inputs, streamer=streamer, **kwargs)

        logger.info("Starting generation thread")
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()

        output = ""
        try:
            for new_text in streamer:
                output += new_text
                if output.endswith("</s>"):
                    output = output[:-4]
                    break
        except Empty:
            logger.warning("Timeout occurred during generation")

        logger.info("Generation completed")
        return output
    except Exception as e:
        logger.error(f"Error in generate_response: {str(e)}")
        return f"Wystąpił błąd: {str(e)}"

def function(prompt, history=[]):
    chat = "<s>"
    for user_prompt, bot_response in history:
        chat += f"[INST] {user_prompt} [/INST] {bot_response}</s> <s>"
    chat += f"[INST] {prompt} [/INST]"
    kwargs = dict(
        max_new_tokens=4096,
        do_sample=True,
        temperature=0.5,
        top_p=0.95,
        repetition_penalty=1.0
    )

    return generate_response(chat, kwargs)

interface = gr.ChatInterface(
    fn=function,
    chatbot=gr.Chatbot(
        avatar_images=None,
        container=False,
        show_copy_button=True,
        layout='bubble',
        render_markdown=True,
        line_breaks=True
    ),
    css='h1 {font-size:22px;} h2 {font-size:20px;} h3 {font-size:18px;} h4 {font-size:16px;}',
    autofocus=True,
    fill_height=True,
    analytics_enabled=False,
    submit_btn='Chat',
    stop_btn=None,
    retry_btn=None,
    undo_btn=None,
    clear_btn=None
)

interface.launch(show_api=True, share=True)