from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np
import random
import json
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel

# Lade NUR RecipeBERT Modell
bert_model_name = "alexdseo/RecipeBERT"
bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
bert_model = AutoModel.from_pretrained(bert_model_name)
bert_model.eval() # Setze das Modell in den Evaluationsmodus

# T5-Modell und -Logik KOMPLETT ENTFERNT für diesen Schritt
# special_tokens und tokens_map sind nicht mehr relevant, bleiben aber als Kommentar

# --- RecipeBERT-spezifische Funktionen ---
def get_embedding(text):
    """Berechnet das Embedding für einen Text mit Mean Pooling über alle Tokens."""
    inputs = bert_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = bert_model(**inputs)
    attention_mask = inputs['attention_mask']
    token_embeddings = outputs.last_hidden_state
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    return (sum_embeddings / sum_mask).squeeze(0)

def average_embedding(embedding_list):
    """Berechnet den Durchschnitt einer Liste von Embeddings."""
    tensors = torch.stack(embedding_list) # embedding_list enthält hier direkt die Tensoren
    return tensors.mean(dim=0)

def get_cosine_similarity(vec1, vec2):
    """Berechnet die Cosinus-Ähnlichkeit zwischen zwei Vektoren."""
    if torch.is_tensor(vec1): vec1 = vec1.detach().numpy()
    if torch.is_tensor(vec2): vec2 = vec2.detach().numpy()
    vec1 = vec1.flatten()
    vec2 = vec2.flatten()
    dot_product = np.dot(vec1, vec2)
    norm_a = np.linalg.norm(vec1)
    norm_b = np.linalg.norm(vec2)
    if norm_a == 0 or norm_b == 0: return 0
    return dot_product / (norm_a * norm_b)


# find_best_ingredients (modifiziert, um die ähnlichste Zutat mit RecipeBERT zu finden)
def find_best_ingredients(required_ingredients, available_ingredients, max_ingredients=6):
    """
    Findet die besten Zutaten: Alle benötigten + EINE ähnlichste aus den verfügbaren Zutaten.
    """
    required_ingredients = list(set(required_ingredients))
    available_ingredients = list(set([i for i in available_ingredients if i not in required_ingredients]))

    final_ingredients = required_ingredients.copy()

    # Nur wenn wir noch Platz haben und zusätzliche Zutaten verfügbar sind
    if len(final_ingredients) < max_ingredients and len(available_ingredients) > 0:
        if final_ingredients:
            # Berechne den Durchschnitts-Embedding der benötigten Zutaten
            required_embeddings = [get_embedding(ing) for ing in required_ingredients]
            avg_required_embedding = average_embedding(required_embeddings)
            
            best_additional_ingredient = None
            highest_similarity = -1.0

            # Finde die ähnlichste Zutat aus den verfügbaren
            for avail_ing in available_ingredients:
                avail_embedding = get_embedding(avail_ing)
                similarity = get_cosine_similarity(avg_required_embedding, avail_embedding)
                if similarity > highest_similarity:
                    highest_similarity = similarity
                    best_additional_ingredient = avail_ing
            
            if best_additional_ingredient:
                final_ingredients.append(best_additional_ingredient)
                print(f"INFO: Added '{best_additional_ingredient}' (similarity: {highest_similarity:.2f}) as most similar.")
        else:
            # Wenn keine benötigten Zutaten, wähle zufällig eine aus den verfügbaren (wie zuvor)
            random_ingredient = random.choice(available_ingredients)
            final_ingredients.append(random_ingredient)
            print(f"INFO: No required ingredients. Added random available ingredient: '{random_ingredient}'.")

    # Begrenze auf max_ingredients, falls durch Zufall/ähnlichster Auswahl zu viele hinzugefügt wurden
    return final_ingredients[:max_ingredients]


# mock_generate_recipe (bleibt gleich)
def mock_generate_recipe(ingredients_list):
    """Generiert ein Mock-Rezept, da T5-Modell entfernt ist."""
    title = f"Einfaches Rezept mit {', '.join(ingredients_list[:3])}" if ingredients_list else "Einfaches Testrezept"
    return {
        "title": title,
        "ingredients": ingredients_list, # Die "generierten" Zutaten sind einfach die Eingabe
        "directions": [
            "Dies ist ein generierter Text von RecipeBERT (ohne T5).",
            "Das Laden des RecipeBERT-Modells war erfolgreich!",
            f"Basierend auf deinen Eingaben wurde '{ingredients_list[-1]}' als ähnlichste Zutat hinzugefügt." if len(ingredients_list) > 1 else "Keine zusätzliche Zutat hinzugefügt."
        ],
        "used_ingredients": ingredients_list # In diesem Mock-Fall sind alle "used"
    }


def process_recipe_request_logic(required_ingredients, available_ingredients, max_ingredients, max_retries):
    """
    Kernlogik zur Verarbeitung einer Rezeptgenerierungsanfrage.
    Für diesen Test wird nur RecipeBERT zum Laden getestet und ein Mock-Rezept zurückgegeben.
    """
    if not required_ingredients and not available_ingredients:
        return {"error": "Keine Zutaten angegeben"}
    try:
        # Hier wird die neue find_best_ingredients verwendet
        optimized_ingredients = find_best_ingredients(
            required_ingredients, available_ingredients, max_ingredients
        )
        
        # Rufe die Mock-Generierungsfunktion auf
        recipe = mock_generate_recipe(optimized_ingredients) 
        
        result = {
            'title': recipe['title'],
            'ingredients': recipe['ingredients'],
            'directions': recipe['directions'],
            'used_ingredients': optimized_ingredients # Jetzt wirklich die vom find_best_ingredients
        }
        return result
    except Exception as e:
        return {"error": f"Fehler bei der Rezeptgenerierung: {str(e)}"}

# --- FastAPI-Implementierung ---
app = FastAPI(title="AI Recipe Generator API (RecipeBERT Only Test)")

class RecipeRequest(BaseModel):
    required_ingredients: list[str] = []
    available_ingredients: list[str] = []
    max_ingredients: int = 7
    max_retries: int = 5 # Wird hier nicht direkt genutzt, aber im Payload beibehalten
    ingredients: list[str] = [] # Für Abwärtskompatibilität

@app.post("/generate_recipe") # Der API-Endpunkt für Flutter
async def generate_recipe_api(request_data: RecipeRequest):
    final_required_ingredients = request_data.required_ingredients
    if not final_required_ingredients and request_data.ingredients:
        final_required_ingredients = request_data.ingredients

    result_dict = process_recipe_request_logic(
        final_required_ingredients,
        request_data.available_ingredients,
        request_data.max_ingredients,
        request_data.max_retries # max_retries wird nur an die Logik übergeben, aber nicht verwendet
    )
    return JSONResponse(content=result_dict)

@app.get("/")
async def read_root():
    return {"message": "AI Recipe Generator API is running (RecipeBERT only, 1 similar ingredient)!"} # Angepasste Nachricht

print("INFO: FastAPI application script finished execution and defined 'app' variable.")