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()

# --- טעינת המודלים ---
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...")
whisper_model = whisper.load_model("large-v2") 
print("Models loaded successfully!")

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

# --- פונקציות עזר ---
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

# --- נתיבי 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")
async def transcribe(file: UploadFile = File(...)):
    print(f"DEBUG: --- FILE RECEIVED: {file.filename}, SIZE: {file.size} ---")
    
    # במקום להשתמש בשם המוזר של הקובץ, ניתן לו שם קבוע ופשוט
    temp_path = f"temp_{uuid.uuid4()}.mp3"
    
    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,
            verbose=False,
            condition_on_previous_text=False
        )
        
        transcription_text = result["text"].strip()
        print(f"DEBUG: --- TRANSCRIPTION RESULT: '{transcription_text}' ---")
        return {"status": "success", "transcription": transcription_text}

    except Exception as e:
        print(f"DEBUG: --- EXCEPTION: {str(e)} ---")
        return {"status": "error", "message": str(e)}

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