import numpy as np
import pandas as pd
import faiss
import torch
import gc
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr

# GPU Bellek yönetimi
def clear_gpu_cache():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()

# Veri yükleme (Space'e yüklediğiniz dosya)
df = pd.read_json("ORNEK_VERILER_160.jsonl", lines=True)

# Metin dönüşüm
def row_to_text(r):
    return (f"Hat {r['hat_no']} | İstasyon {r['istasyon_no']} | Model: {r['model']} | "
            f"CPU: {r['cpu_model']} | RAM: {r['ram']} GB | Disk: {r['disk']} GB {r['depolama_turu']} | "
            f"GPU: {r['gpu_model']} ({r['gpu_gb']} GB) | Uygulama: {r['kullanilan_program']} | "
            f"CPU%: {r['cpu_kullanimi']}, GPU%: {r['gpu_kullanimi']}, RAM%: {r['ram_kullanimi']}")

docs = df.apply(row_to_text, axis=1).tolist()

# Embedding modeli
emb_model = SentenceTransformer("intfloat/multilingual-e5-small")

def embed(texts):
    embeddings = emb_model.encode(texts, normalize_embeddings=True, convert_to_numpy=True)
    return np.ascontiguousarray(embeddings.astype(np.float32))

# FAISS indeksi
embs = embed(docs)
index = faiss.IndexFlatIP(embs.shape[1])
index.add(embs)

# LLM modeli
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=dtype)

# Yardımcı fonksiyonlar
def retrieve(query, k=8):
    q_emb = embed([query])
    D, I = index.search(q_emb, k)
    ctx_rows = [docs[i] for i in I[0]]
    ctx_df = df.iloc[I[0]].copy()
    return ctx_rows, ctx_df

def risk_and_bottleneck(cpu, gpu, ram):
    risk = 0.35*cpu + 0.30*ram + 0.25*gpu
    tags = []
    if cpu >= 85 and gpu < 70: tags.append("CPU")
    if gpu >= 85 and cpu < 75: tags.append("GPU")
    if ram >= 85: tags.append("RAM")
    return round(risk, 3), ",".join(tags) if tags else "YOK"

def tool_answer_if_detected(q, ctx_df):
    ql = q.lower()
    if any(k in ql for k in ["risk", "darboğaz", "darbogaz", "bottleneck"]):
        out = []
        for _, r in ctx_df.iterrows():
            risk, bot = risk_and_bottleneck(r["cpu_kullanimi"], r["gpu_kullanimi"], r["ram_kullanimi"])
            out.append(f"{r['istasyon_no']} ({r['model']}, {r['kullanilan_program']}): risk={risk}, darbogaz={bot}")
        return "• " + "\n• ".join(out)
    return None

SYSTEM_PROMPT = "Sen teknik bir destek asistanısın. Bağlamdaki satırlara dayanarak soruya net, Türkçe ve adım adım cevap ver. Varsayım yapma; bağlamda yoksa açıkça söyle."

def chat_once(question, k=8, max_new_tokens=420):
    clear_gpu_cache()
    ctx_rows, ctx_df = retrieve(question, k=k)
    tool = tool_answer_if_detected(question, ctx_df)
    context_block = "\n".join(f"- {c}" for c in ctx_rows)
    if tool:
        context_block += "\n\n[HESAP]:\n" + tool
        
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Soru: {question}\n\nBağlam:\n{context_block}"}
    ]
    
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, temperature=0.2, top_p=0.9, pad_token_id=tokenizer.eos_token_id)
    
    text = tokenizer.decode(out[0], skip_special_tokens=True)
    answer = text.split("assistant\n")[-1].strip()
    return answer

# Gradio arayüzü
def respond(message, history):
    try:
        return chat_once(message)
    except Exception as e:
        return f"Hata: {str(e)}\n\nLütfen soruyu yeniden formüle edin."

demo = gr.ChatInterface(
    fn=respond,
    title="PC Performans Asistanı (RAG + Mini LLM)",
    description="Örnek: 'En riskli 5 bilgisayar', 'SolidWorks için uygun PC var mı?', 'w127_709 istasyonunun durumu?'",
    examples=[
        "En riskli 3 bilgisayarı listeler misin?",
        "SolidWorks için uygun PC var mı?",
        "w127_709 istasyonunun durumu nedir?",
        "RAM kullanımı yüksek olan istasyonlar hangileri?"
    ]
)

demo.launch()