import os
import time
import json
import shutil
from pathlib import Path
from dotenv import load_dotenv
from google import genai
from PIL import Image

# 1. Konfiguration & Umgebung
load_dotenv()
API_KEY = os.getenv("GEMINI_API_KEY")

# Verzeichnisse definieren
BASE_DIR = Path(__file__).parent.absolute()
SRC_DIR = BASE_DIR / "img_new"
DEST_DIR = BASE_DIR / "img_erl"
ERROR_DIR = BASE_DIR / "img_error"
JSON_DIR = BASE_DIR / "json"

# Ordner sicherstellen
for d in [DEST_DIR, ERROR_DIR, JSON_DIR]:
    d.mkdir(exist_ok=True, parents=True)

# Gemini Client
client = genai.Client(api_key=API_KEY)
MODEL_ID = "models/gemini-3-flash-preview"

PROMPT = """Du bist ein spezialisierter Assistent zur Datenextraktion aus Bilddateien. Deine Aufgabe ist es, Informationen von einer Rezeptkarte in ein präzises JSON-Format zu überführen.

Analyse der Bildbereiche:
Rezeptname: Befindet sich oben links in der weißen Box.
Rezeptgruppe: Befindet sich oben rechts in der weißen Box.
QR-Code: Befindet sich unten in der Mitte.
Zutaten: Liste auf der rechten Seite unter der Überschrift "Zutaten".
Optionale Zutaten: Falls vorhanden, unter der Überschrift "Optional".
Gewürze: Falls vorhanden, unter der entsprechenden Überschrift.

JSON-Strukturvorgaben:
Verwende exakt die folgende Struktur. Wenn eine Kategorie fehlt, setze null oder [].
Achte darauf, Menge und Einheit korrekt zu trennen. Falls keine Einheit angegeben ist, verwende null.

Wichtig: Gib als Antwort ausschließlich den reinen JSON-Code zurück. Keine Erklärungen, kein Markdown-Code-Block.

{
  "rezeptname": "String oder null",
  "rezeptgruppe": "String oder null",
  "qrcode": "String oder null",
  "zutaten": [],
  "optional": [],
  "gewuerze": []
}"""

def process_agent():
    valid_extensions = ('.jpg', '.jpeg', '.png')
    images = sorted([f for f in SRC_DIR.iterdir() if f.suffix.lower() in valid_extensions])
    
    if not images:
        print("Warten auf neue Bilder in 'img_new'...")
        return

    for img_path in images:
        print(f"\n--- Verarbeitung: {img_path.name} ---")
        success = False
        
        try:
            # KI-Analyse
            pil_img = Image.open(img_path)
            response = client.models.generate_content(
                model=MODEL_ID,
                contents=[PROMPT, pil_img]
            )
            
            # JSON-Bereinigung
            raw_text = response.text.strip()
            if "```" in raw_text:
                raw_text = raw_text.split("```")[1].replace("json", "").strip()
            
            recipe_data = json.loads(raw_text)

            # Validierung & Speichern
            if recipe_data.get("rezeptname") and recipe_data.get("rezeptgruppe"):
                json_file = JSON_DIR / f"{img_path.stem}.json"
                with open(json_file, "w", encoding="utf-8") as f:
                    json.dump(recipe_data, f, ensure_ascii=False, indent=2)
                
                shutil.move(str(img_path), str(DEST_DIR / img_path.name))
                print(f"ERFOLG: {json_file.name} gespeichert.")
                success = True
            else:
                missing = []
                if not recipe_data.get("rezeptname"): missing.append("Name")
                if not recipe_data.get("rezeptgruppe"): missing.append("Gruppe")
                print(f"ABBRUCH: Fehlende Daten ({', '.join(missing)})")

        except Exception as e:
            print(f"FEHLER: {e}")

        # Bei Fehlern oder unvollständigen Daten verschieben
        if not success:
            shutil.move(str(img_path), str(ERROR_DIR / img_path.name))
            print("Bild in den Error-Ordner verschoben.")

        # Cooldown
        time.sleep(5)

if __name__ == "__main__":
    process_agent()