import torch
import os
import shutil
import whisper
import uuid
from itertools import product
from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from typing import List
from transformers import AutoTokenizer, AutoModelForMaskedLM

app = FastAPI()

# --- 1. טעינת המודלים (ייטענו פעם אחת כשהשרת עולה) ---
print("Loading DictaBERT model...")
model_path = "dicta-il/dictabert"
tokenizer = AutoTokenizer.from_pretrained(model_path)
bert_model = AutoModelForMaskedLM.from_pretrained(model_path)
bert_model.eval()

print("Loading Whisper model (this might take a minute)...")
# טוענים את המודל הגדול. השרת החינמי הוא CPU, אז הוא ירוץ עליו.
whisper_model = whisper.load_model("large-v2") 
print("Models loaded successfully!")

# --- 2. מבני נתונים ---
class PredictionRequest(BaseModel):
    template: str
    lists: List[List[str]]

# --- 3. לוגיקה ---
def get_top_k_combinations(template, lists_of_candidates, k=3):
    num_masks = len(lists_of_candidates)
    text = template
    for i in range(1, num_masks + 1):
        text = text.replace(f"{{m{i}}}", tokenizer.mask_token)

    inputs = tokenizer(text, return_tensors="pt")
    mask_indices = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]

    with torch.no_grad():
        outputs = bert_model(**inputs)
        logits = outputs.logits

    all_combinations_with_scores = []
    for combination in product(*lists_of_candidates):
        current_combined_score = 0
        for i, word in enumerate(combination):
            word_id = tokenizer.convert_tokens_to_ids(word)
            current_combined_score += logits[0, mask_indices[i], word_id].item()
        all_combinations_with_scores.append((combination, current_combined_score))

    return sorted(all_combinations_with_scores, key=lambda x: x[1], reverse=True)[:k]

def format_final_sentence(template, best_comb):
    sentence = template
    for i, word in enumerate(best_comb):
        target = "{m" + str(i + 1) + "}"
        sentence = sentence.replace(target, word)
    return sentence

# --- 4. נתיבי API ---
@app.get("/")
def home():
    return {"message": "Speechly API is running!"}

@app.post("/predict")
async def predict(data: PredictionRequest):
    top_combinations = get_top_k_combinations(data.template, data.lists, k=3)
    final_sentences = [format_final_sentence(data.template, c[0]) for c in top_combinations]
    top_words = [list(c[0]) for c in top_combinations]
    return {"full_sentences": final_sentences, "mask_selections": top_words}

@app.post("/transcribe")
def transcribe(file: UploadFile = File(...)):
    temp_path = f"temp_{uuid.uuid4()}_{file.filename}"
    try:
        with open(temp_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        result = whisper_model.transcribe(
            temp_path,
            language="he",
            temperature=0,
            fp16=False, # אנחנו על CPU אז זה חובה להיות False
            verbose=False,
            condition_on_previous_text=False
        )
        return {"status": "success", "transcription": result["text"].strip()}

    except Exception as e:
        return {"status": "error", "message": str(e)}

    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)