import os
import cv2
import numpy as np
from fastapi import FastAPI, Request, Response, Form
from ultralytics import YOLO
from datetime import datetime
import psycopg2

app = FastAPI()
model = YOLO("best.pt")

# DB CONFIG
DB_CONFIG = {
    'user': 'postgres',
    'host': 'localhost',
    'database': 'chatiere_db',
    'password': 'sijereste',
    'port': '5432'
}

PHOTOS_HISTORIQUE = "../photos_historique"
PHOTOS_ANALYSE = "photos_analyse"

# Variables globales pour suivre la session en cours
compteur_photo = 1
nom_photo_historique = ""
photo_historique = None
current_chat_id = None  # Stockera l'ID de la DB du chat qui vient de badger

def enregistrer_passage(resultat_proie, chemin_photo):
    global current_chat_id
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()
        maintenant = datetime.now()

        # On utilise current_chat_id récupéré lors du badgeage
        # Si pour une raison X il est vide, on met 0 ou NULL
        id_a_enregistrer = current_chat_id if current_chat_id else None

        sql = """
            INSERT INTO passages (chat_id, passage_date, proie, photo)
            VALUES (%s, %s, %s, %s)
        """
        cursor.execute(sql, (id_a_enregistrer, maintenant, resultat_proie, chemin_photo))

        conn.commit()
        cursor.close()
        conn.close()
        print(f"Enregistré en DB pour le chat {id_a_enregistrer} (Proie: {resultat_proie})")
    except Exception as e:
        print(f"Erreur PostgreSQL enregistrer_passage : {e}")

def vider_dossier_analyse():
    if not os.path.exists(PHOTOS_ANALYSE):
        os.makedirs(PHOTOS_ANALYSE)
    for filename in os.listdir(PHOTOS_ANALYSE):
        file_path = os.path.join(PHOTOS_ANALYSE, filename)
        try:
            if os.path.isfile(file_path):
                os.remove(file_path)
        except Exception as e:
            print(f"Erreur nettoyage : {e}")

# ========================================================
# NOUVELLE ROUTE : VERIFICATION DU BADGE
# ========================================================
@app.post("/verifier_puce")
async def verifier_badge(id: str = Form(...)):
    global current_chat_id
    print(f"Vérification DB pour le RFID : {id}")
    
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        cursor = conn.cursor()
        
        # Chercher le nom et l'ID du chat
        cursor.execute("SELECT id, nom FROM chats WHERE puce_numero = %s AND chatiere_id = 1", (id,))
        row = cursor.fetchone()
        
        cursor.close()
        conn.close()

        if row:
            current_chat_id = row[0] # On stocke l'ID pour le log photo plus tard
            nom_chat = row[1]
            print(f"Chat reconnu : {nom_chat} (ID: {current_chat_id})")
            return Response(content=nom_chat) # Envoie le nom à l'ESP32
        else:
            current_chat_id = None
            print("Badge inconnu !")
            return Response(content="NON")
            
    except Exception as e:
        print(f"Erreur DB verifier_badge : {e}")
        return Response(content="ERROR")

# ========================================================
# ROUTE ANALYSE PHOTO (Inchangée mais utilise current_chat_id)
# ========================================================
@app.post("/analyse_photo")
async def image_reception(request: Request):
    global compteur_photo, nom_photo_historique, photo_historique

    if compteur_photo == 1:
        vider_dossier_analyse()
        nom_photo_historique = ""
        photo_historique = None

    data = await request.body()
    file_bytes = np.frombuffer(data, np.uint8)
    image_esp = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    
    if image_esp is None:
        return Response(content="Erreur image")

    # Correction rotation
    image_esp = cv2.rotate(image_esp, cv2.ROTATE_90_CLOCKWISE)

    print(f"Analyse photo n°{compteur_photo}...")

    resultats = model(image_esp, conf=0.5)
    image_annotee = resultats[0].plot()

    cv2.imwrite(f"{PHOTOS_ANALYSE}/photo_{compteur_photo}.jpg", image_annotee)

    souris_detectee = False
    tete_chat_detectee = False
    date = datetime.now().strftime("%Y%m%d_%H%M%S")

    for resultat in resultats:
        for box in resultat.boxes:
            class_name = model.names[int(box.cls[0])]
            if class_name == "prey":
                souris_detectee = True
                nom_photo_historique = f"{PHOTOS_HISTORIQUE}/passage_{date}_PROIE.jpg"
                photo_historique = image_annotee
            elif class_name == "cat_head":
                tete_chat_detectee = True
                if not souris_detectee:
                    nom_photo_historique = f"{PHOTOS_HISTORIQUE}/passage_{date}_OK.jpg"
                    photo_historique = image_annotee

    if souris_detectee:
        print("SOURIS DÉTECTÉE !")
        cv2.imwrite(nom_photo_historique, photo_historique)
        enregistrer_passage(True, nom_photo_historique)
        compteur_photo = 1 
        return Response(content="MOUSE")

    if compteur_photo < 5:
        compteur_photo += 1
        return Response(content="NEXT")
    else:
        # Fin de la rafale
        compteur_photo = 1
        if tete_chat_detectee:
            print("Chat autorisé sans proie.")
            cv2.imwrite(nom_photo_historique, photo_historique)
            enregistrer_passage(False, nom_photo_historique)
            return Response(content="CAT")
        else:
            print("Aucun chat identifié sur les photos.")
            return Response(content="CLOSE")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)
