import torch
import json
import re
import wikipedia
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from labels_config import CUSTOM_LABELS
from prompts import SYSTEM_COT_PROMPT
import spaces

wikipedia.set_lang("tr")

# Modellerin çalışma anında yüklenmesi için yardımcı fonksiyon
def get_models():
    model_name = "Qwen/Qwen2.5-1.5B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto")
    # NER modeli CPU'da başlayabilir, GPU işlemi sırasında taşınır
    ner = pipeline("ner", model="xlm-roberta-large-finetuned-conll03-english", aggregation_strategy="simple")
    return tokenizer, model, ner

def get_wiki_summary(query):
    try:
        search_results = wikipedia.search(query)
        if not search_results: return None
        return wikipedia.summary(search_results[0], sentences=2, auto_suggest=False)
    except: return None

# CUDA hatasını engellemek için ana NER işlemini de bu kapsama alabiliriz
def ner_pipe(text):
    _, _, ner = get_models()
    return ner(text)

@spaces.GPU
def refine_labels_batch(misc_items, full_sentence):
    if not misc_items: return []
    tokenizer, llm_model, _ = get_models()
    
    label_rules = ""
    for k, v in CUSTOM_LABELS.items():
        label_rules += f"### {k}\nTANIM: {v['tanim']}\nANAHTARLAR: {', '.join(v['anahtar_kelimeler'])}\n\n"
    
    targets_text = "".join([f"- VARLIK: {item['word']} | WIKI: {item['wiki']}\n" for item in misc_items])
    final_prompt = SYSTEM_COT_PROMPT.format(label_rules=label_rules, full_sentence=full_sentence, targets=targets_text)
    
    messages = [{"role": "user", "content": final_prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(llm_model.device)
    
    try:
        generated_ids = llm_model.generate(model_inputs.input_ids, max_new_tokens=1000, do_sample=False)
        output = tokenizer.batch_decode(generated_ids[:, model_inputs.input_ids.shape[1]:], skip_special_tokens=True)[0]
        
        json_match = re.search(r'\[\s*\{.*\}\s*\]', output, re.DOTALL)
        if json_match:
            raw_json = json_match.group(0).strip()
            if raw_json.count('{') > raw_json.count('}'): raw_json += "}]"
            results = json.loads(raw_json)
            # app.py'de KeyError oluşmaması için anahtarları garanti altına alıyoruz
            for r in results:
                r.setdefault('varlik', 'Bilinmeyen')
                r.setdefault('karar', 'MISC')
                r.setdefault('reasoning', 'Analiz yapılamadı.')
            return results
        return []
    except Exception as e:
        print(f"Hata: {e}")
        return []