import torch, traceback
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
from log import log
from pydantic import BaseModel

model = None
tokenizer = None
eos_token_id = None

class Message(BaseModel):
    user_input: str

def setup_model(s_config):
    global model, tokenizer, eos_token_id
    try:
        log("🧠 setup_model() başladı")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        log(f"📡 Kullanılan cihaz: {device}")
        tokenizer = AutoTokenizer.from_pretrained(s_config.MODEL_BASE, use_fast=False)
        log("📦 Tokenizer yüklendi. Ana model indiriliyor...")
        model = AutoModelForCausalLM.from_pretrained(s_config.MODEL_BASE, torch_dtype=torch.float32).to(device)
        log("📦 Ana model indirildi ve yüklendi. eval() çağırılıyor...")
        tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
        model.config.pad_token_id = tokenizer.pad_token_id
        eos_token_id = tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"][0]
        model.eval()
        log("✅ Ana model eval() çağrıldı")
        log(f"📦 Intent modeli indiriliyor: {s_config.INTENT_MODEL_ID}")
        _ = AutoTokenizer.from_pretrained(s_config.INTENT_MODEL_ID)
        _ = AutoModelForSequenceClassification.from_pretrained(s_config.INTENT_MODEL_ID)
        log("✅ Intent modeli önbelleğe alındı.")
        log("✔️ Model başarıyla yüklendi ve sohbet için hazır.")
    except Exception as e:
        log(f"❌ setup_model() hatası: {e}")
        traceback.print_exc()

async def generate_response(text, app_config):
    messages = [{"role": "user", "content": text}]
    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
    eos_token = tokenizer("<|im_end|>", add_special_tokens=False)["input_ids"][0]
    input_ids = encodeds.to(model.device)
    attention_mask = (input_ids != tokenizer.pad_token_id).long()

    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=128,
            do_sample=app_config.USE_SAMPLING,
            eos_token_id=eos_token,
            pad_token_id=tokenizer.pad_token_id,
            return_dict_in_generate=True,
            output_scores=True
        )

    if not app_config.USE_SAMPLING:
        scores = torch.stack(output.scores, dim=1)
        probs = torch.nn.functional.softmax(scores[0], dim=-1)
        top_conf = probs.max().item()
    else:
        top_conf = None

    decoded = tokenizer.decode(output.sequences[0], skip_special_tokens=True).strip()
    for tag in ["assistant", "<|im_start|>assistant"]:
        start = decoded.find(tag)
        if start != -1:
            decoded = decoded[start + len(tag):].strip()
            break
    return decoded, top_conf