from fastapi import FastAPI
from pydantic import BaseModel
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
import os
import google.generativeai as genai
import pandas as pd
import numpy as np

# === 1. Configurer le cache dans /tmp/ (accessible en écriture) ===
os.environ['HF_HOME'] = '/tmp/huggingface'
os.environ['TRANSFORMERS_CACHE'] = '/tmp/huggingface'
os.makedirs('/tmp/huggingface', exist_ok=True)

# === 2. Charger le modèle avec gestion des erreurs ===
try:
    model = SentenceTransformer('all-MiniLM-L6-v2', cache_folder='/tmp/huggingface')
except Exception as e:
    print(f"Erreur de chargement (nom court): {e}")
    model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2', cache_folder='/tmp/huggingface')

# === 3. Configurer ChromaDB dans /tmp/ ===
client = chromadb.PersistentClient(path="/tmp/chroma_db")

# === 4. Configurer Gemini ===
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
genai.configure(api_key=GEMINI_API_KEY)
model_gemini = genai.GenerativeModel("gemini-2.0-flash-lite")

app = FastAPI()

# === Classe Query pour FastAPI ===
class Query(BaseModel):
    query: str

# === Fonction d'initialisation de la collection ===
def init_collection():
    # Essaie de récupérer la collection, sinon la crée et l'initialise
    try:
        collection = client.get_collection("Festivals")
        print("Collection 'Festivals' chargée")
    except Exception:
        print("Collection 'Festivals' non trouvée, création et remplissage...")
        collection = client.create_collection(name="Festivals", metadata={"hnsw:space": "ip"})
        
        # Ici tu dois récupérer tes données (ex: depuis Airtable)
        # Exemple fictif avec dataframe statique, remplace par ta vraie fonction fetch_data()
        df = pd.DataFrame([
            {
                "Festival Name": "Hellfest",
                "City": "Clisson",
                "Country": "France",
                "Dates": "15-18 Juin",
                "Genre": "Metal/Rock",
                "Ticket Price (EUR)": 200,
                "Accommodation option": "Camping",
                "Atmosphere": "Festive"
            },
            {
                "Festival Name": "Sziget",
                "City": "Budapest",
                "Country": "Hongrie",
                "Dates": "10-15 Août",
                "Genre": "Divers",
                "Ticket Price (EUR)": 150,
                "Accommodation option": "Tentes",
                "Atmosphere": "Jeune et dynamique"
            }
        ])

        embeddings, metadatas, ids = [], [], []
        for idx, row in df.iterrows():
            row_text = "\n".join([f"{col}: {val}" for col, val in row.items() if pd.notna(val)])
            embedding = model.encode(row_text)
            embedding /= np.linalg.norm(embedding)
            embeddings.append(embedding.tolist())
            metadatas.append({k: str(v) for k, v in row.items() if pd.notna(v)})
            ids.append(f"fest_{idx}")

        collection.add(embeddings=embeddings, metadatas=metadatas, ids=ids)
        print(f"Collection 'Festivals' initialisée avec {len(ids)} entrées")

    return collection


# === Initialisation au démarrage ===
collection = init_collection()


@app.get("/")
async def root():
    return {"message": "🎵 GrooveNomad Festivals API est en ligne !"}

@app.post("/api/ask")
async def ask_festival(data: Query):
    query_embedding = model.encode(data.query).tolist()
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=3,
        include=["metadatas", "distances"]
    )

    def get_val(meta, key):
        val = meta.get(key)
        if not val or str(val).lower() in ["", "nan", "none"]:
            return "Non spécifié"
        return val

    context = "\n".join([
        f"{get_val(m, 'Festival Name')} à {get_val(m, 'City')} ({get_val(m, 'Country')}) le {get_val(m, 'Dates')} 🎶 {get_val(m, 'Genre')} – {get_val(m, 'Ticket Price (EUR)')}€"
        for m in results["metadatas"][0]
    ])

    prompt = f"""Tu es un expert des festivals. Voici la requête :
\"{data.query}\"
Voici les options :
{context}
Génère une réponse naturelle, enthousiaste, concise, avec des emojis."""

    try:
        response = model_gemini.generate_content(prompt)
        return {"response": response.text}
    except Exception as e:
        print(f"Erreur lors de la génération : {e}")
        return {"error": str(e)}