import streamlit as st
from transformers import LlamaTokenizerFast
from peft import AutoPeftModelForCausalLM
import torch
from io import BytesIO


# ------------------------------------------------
# Model ve tokenizer'ı sadece bir kez yükleyelim:
# ------------------------------------------------
@st.cache_resource
def load_model_and_tokenizer():
    tokenizer = LlamaTokenizerFast.from_pretrained("anilguleroglu/cllm-0.0.2-chat-final")
    model = AutoPeftModelForCausalLM.from_pretrained("anilguleroglu/cllm-0.0.2-chat-final")
    return model, tokenizer

model, tokenizer = load_model_and_tokenizer()


# ------------------------------------------------
# Metin üretim fonksiyonu
# ------------------------------------------------
def generate_text(prompt, max_length=1024):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    outputs = model.generate(
        inputs.input_ids,
        attention_mask=inputs.attention_mask,
        max_length=max_length,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        top_k=10,
        top_p=0.9,
        temperature=0.3,
        repetition_penalty=1.3
    )

    response =  tokenizer.decode(outputs[0], skip_special_tokens=False)
    print(len(response))
    keyword = "<|start_header_id|> assistant<|end_header_id|>"
    print("RESPONSE:")
    print(response)
    # Anahtar kelimeden sonrasını çıkarma
    if keyword in response:
        extracted_content = response.split(keyword, 1)[1].strip().replace("<|eot_id|>","")
        return extracted_content
    else:
        response
    



# ------------------------------------------------
# Chat için prompt oluşturma fonksiyonu
# (system mesajını eklemek isteğe bağlı; 
#  bu örnekte system mesajı "apply_chat_template" 
#  fonksiyonuna elle eklenebilir.)
# ------------------------------------------------
def build_prompt(system_message, chat_history):
    """
    system_message: string
    chat_history:   [
        {"role":"user", "content":"..."},
        {"role":"assistant", "content":"..."},
        ...
    ]
    """
    # System mesajını chat'in başına eklemek için:
    # (İsterseniz LlamaTokenizerFast'in 'apply_chat_template'
    # metodunu da kullanabilirsiniz.)
    
    # Bu örnekte basit bir şekilde kendimiz birleştiriyoruz:
    chat_history = [{"role": "system", "content": system_message}] + chat_history
    return tokenizer.apply_chat_template(chat_history, tokenize=False, add_generation_prompt=True)


# ------------------------------------------------
# Streamlit Arayüzü
# ------------------------------------------------
def main():
    st.title("CLLM-0.0.2 Chat")

    # Session state'de system mesajı ve konuşmaları tutalım
    if "system_message" not in st.session_state:
        st.session_state["system_message"] = ""

    if "messages" not in st.session_state:
        st.session_state["messages"] = []  # {"role":"user"/"assistant", "content":...}

    # System mesajı giriş alanı
    st.session_state["system_message"] = st.text_area(
        "System Mesajı (konuşmayı resetleseniz de kaybolmaz)",
        st.session_state["system_message"],
        height=100
    )

    # Reset buttonu
    if st.button("Konuşmayı Sıfırla"):
        st.session_state["messages"] = []

    # Mevcut konuşmayı ekranda göster
    st.subheader("Konuşma Geçmişi")
    for msg in st.session_state["messages"]:
        if msg["role"] == "user":
            with st.chat_message("user"):
                st.write(msg["content"])
        else:
            with st.chat_message("assistant"):
                st.write(msg["content"])

    # Kullanıcı girişi
    user_input = st.chat_input(placeholder="Mesajınızı yazın ve Enter'a basın...")
    if user_input:
        # 1) Kullanıcı mesajını ekle
        st.session_state["messages"].append({"role": "user", "content": user_input})

        # 2) Prompt'u oluştur
        prompt = build_prompt(st.session_state["system_message"], st.session_state["messages"])
        
        # 3) Modelden cevap al
        with st.chat_message("assistant"):
            with st.spinner("Model cevap üretiyor..."):
                generated_text = generate_text(prompt)

    
        # 5) Asistan cevabını kaydet
        st.session_state["messages"].append({"role": "assistant", "content": generated_text})

        # 6) Ekrana yaz
        st.write(generated_text)
    
    # PDF Download
    st.divider()


if __name__ == "__main__":
    main()