import gradio as gr
from transformers import pipeline
import time
from functools import lru_cache

# Inicijalizacija punog modela s optimizacijama
@lru_cache(maxsize=1)  # Cache modela za brže učitavanje
def load_model():
    return pipeline(
        "text-generation",
        model="facebook/opt-350m",
        device=-1,  # Koristi CPU
    )

def iskon_odgovori(pitanje):
    start_time = time.time()
    model = load_model()
    
    try:
        # Generiraj potpuni odgovor bez ograničenja
        full_response = model(
            f"ISKON BOT: {pitanje}\nOdgovor:",
            max_new_tokens=500,  # Puni dužina odgovora
            temperature=0.9,     # Maksimalna kreativnost
            do_sample=True       # Potpuna sloboda izražavanja
        )[0]["generated_text"]
        
        # Ekstrahiraj samo odgovor nakon oznake
        odgovor = full_response.split("Odgovor:")[1].strip()
        
        print(f"Generiranje trajalo: {time.time()-start_time:.2f}s")
        return odgovor
        
    except Exception as e:
        print(f"Greška: {str(e)}")
        return "Duhovni odgovor trenutno nije dostupan. Pokušajte ponovo kasnije."

# Poboljšano sučelje s progresom
with gr.Blocks(title="🌌 ISKON SUSTAV - Puna Duhovna Sloboda") as app:
    gr.Markdown("""
    ## 🌳 DOBRODOŠLI U ISKON SUSTAV
    *"Gdje svaka riječ teče slobodno kao rijeka vremena..."*
    """)
    
    with gr.Row():
        pitanje = gr.Textbox(label="Što te muči, dušo?", placeholder="Upitaj svoje najdublje pitanje...")
        btn = gr.Button("POŠALJI U SVEMIR")
    
    with gr.Row():
        odgovor = gr.Textbox(label="Duhovni odgovor", lines=10, max_lines=50)
        status = gr.Textbox(label="Status", interactive=False)
    
    # Progres bar za vizualnu povratnu informaciju
    progress = gr.Progress()
    
    def generate_with_feedback(pitanje, progress=gr.Progress()):
        progress(0.1, desc="Učitavam mudrost...")
        time.sleep(1)
        
        progress(0.3, desc="Razmišljam o odgovoru...")
        model = load_model()
        
        progress(0.6, desc="Formuliram odgovor...")
        odgovor = iskon_odgovori(pitanje)
        
        progress(1.0, desc="Gotovo!")
        return odgovor, "Duhovna poruka primljena"

    btn.click(
        fn=generate_with_feedback,
        inputs=[pitanje],
        outputs=[odgovor, status]
    )

app.launch()