# ai_gemini_parser.py
import os
import re
import json
import logging
from pathlib import Path
from typing import Any, Dict, Optional, List

try:
    import google.generativeai as genai
except ImportError:
    genai = None

try:
    # Diese Bibliothek wird jetzt benötigt! Installation: pip install PyMuPDF
    import fitz  # PyMuPDF
except ImportError:
    fitz = None

DEFAULT_MODEL = "models/gemini-1.5-flash"
PAGE_BATCH_SIZE = 8  # Anzahl Seiten pro AI-Aufruf, um Limits zu vermeiden

def _configure_model(model_name: str) -> Optional['genai.GenerativeModel']:
    if genai is None:
        raise RuntimeError("google-generativeai ist nicht installiert.")
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        raise RuntimeError("GOOGLE_API_KEY ist nicht als Umgebungsvariable gesetzt.")
    genai.configure(api_key=api_key)
    safety_settings = [
        {"category": c, "threshold": "BLOCK_NONE"}
        for c in ["HARM_CATEGORY_HARASSMENT", "HARM_CATEGORY_HATE_SPEECH", 
                  "HARM_CATEGORY_SEXUALLY_EXPLICIT", "HARM_CATEGORY_DANGEROUS_CONTENT"]
    ]
    return genai.GenerativeModel(model_name, safety_settings=safety_settings)

def _pdf_to_image_parts(pdf_path: Path) -> List[Dict[str, Any]]:
    """Konvertiert PDF-Seiten in Bild-Daten für die Gemini API."""
    if fitz is None:
        raise RuntimeError("PyMuPDF ist nicht installiert (pip install PyMuPDF). PDF-Verarbeitung nicht möglich.")
    
    image_parts = []
    doc = fitz.open(pdf_path)
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        pix = page.get_pixmap(dpi=200)
        image_parts.append({
            "mime_type": "image/png",
            "data": pix.tobytes("png")
        })
    return image_parts

def _call_gemini_with_image_batches(image_parts: List, prompt_text: str, model_name: str) -> Dict:
    """Ruft Gemini in Batches auf, um grosse Dokumente zu verarbeiten."""
    model = _configure_model(model_name)
    if not model:
        raise RuntimeError("Gemini-Modell konnte nicht konfiguriert werden.")
    
    aggregated_results = {}
    
    for i in range(0, len(image_parts), PAGE_BATCH_SIZE):
        batch = image_parts[i:i + PAGE_BATCH_SIZE]
        logging.info(f"Verarbeite PDF-Seiten-Batch {i//PAGE_BATCH_SIZE + 1}...")
        
        full_prompt = [prompt_text] + batch
        
        try:
            response = model.generate_content(
                full_prompt,
                generation_config={"response_mime_type": "application/json"}
            )
            if not response.parts:
                logging.warning(f"Leere Antwort vom Modell für Batch {i//PAGE_BATCH_SIZE + 1}. Finish Reason: {response.candidates[0].finish_reason}")
                continue
            
            batch_result = json.loads(response.text)
            if isinstance(batch_result, dict):
                aggregated_results.update(batch_result)
        except Exception as e:
            logging.error(f"Fehler bei der Gemini-Analyse für Batch {i//PAGE_BATCH_SIZE + 1}: {e}")
            continue
            
    return aggregated_results


def extract_offer_from_pdf(pdf_path: Path, lv_posmap: Optional[Dict[str, Any]] = None, model_name: str = DEFAULT_MODEL) -> Dict[str, Dict[str, Optional[float]]]:
    """Liest Preise (EP/GP) aus einer Offerte (PDF) mit der robusten Batch-Methode."""
    if not lv_posmap:
        pos_list_str = "Bitte extrahieren Sie alle gefundenen NPK-Positionen mit Preisen."
    else:
        pos_list_str = ", ".join(f'"{p}"' for p in lv_posmap.keys())

    prompt = (
        "Du bist ein Experte für Schweizer NPK-Offerten. Analysiere die angehängten Bilder von PDF-Seiten.\n"
        f"Extrahiere NUR für die folgenden Positionsnummern den Einheitspreis (ep) und/oder den Gesamtpreis (gp): [{pos_list_str}]\n"
        "Antworte NUR mit einem JSON-Objekt im Format: "
        '{\n  "<positionsnummer>": {"ep": <zahl_oder_null>, "gp": <zahl_oder_null>}\n}'
    )
    
    try:
        image_parts = _pdf_to_image_parts(pdf_path)
        raw_result = _call_gemini_with_image_batches(image_parts, prompt, model_name)
        
        final_result = {}
        if isinstance(raw_result, dict):
            for pos, data in raw_result.items():
                if isinstance(data, dict):
                    final_result[str(pos)] = {"ep": _num(data.get("ep")), "gp": _num(data.get("gp"))}
        
        if lv_posmap:
            for pos in lv_posmap:
                if pos not in final_result:
                    final_result[pos] = {"ep": None, "gp": None}
        return final_result
    except Exception as ex:
        logging.error(f"Komplette Gemini-Extraktion fehlgeschlagen: {ex}")
        return {pos: {"ep": None, "gp": None} for pos in (lv_posmap or {})}

def extract_lv_from_pdf(pdf_path: Path, model_name: str = DEFAULT_MODEL) -> Dict[str, Dict[str, Any]]:
    """Extrahiert LV-Struktur aus einem PDF mit der robusten Batch-Methode."""
    prompt = (
        "Du bist ein Experte für Schweizer NPK-Leistungsverzeichnisse. Analysiere die angehängten Bilder von PDF-Seiten.\n"
        "Extrahiere alle LV-Positionen mit Menge (qty), Einheit (unit) und Kurztext (text).\n"
        "Antworte NUR mit einem JSON-Objekt im Format:\n"
        '{\n  "<positionsnummer>": {"qty": <zahl_oder_null>, "unit": "<text_oder_null>", "text": "<kurztext>"}\n}'
    )
    try:
        image_parts = _pdf_to_image_parts(pdf_path)
        raw_result = _call_gemini_with_image_batches(image_parts, prompt, model_name)
        
        final_result = {}
        if isinstance(raw_result, dict):
             for pos, data in raw_result.items():
                if isinstance(data, dict):
                    final_result[str(pos)] = {"qty": _num(data.get("qty")), "unit": str(data.get("unit") or "").strip(), "text": str(data.get("text") or "").strip()}
        return final_result
    except Exception as e:
        logging.error(f"LV-Extraktion aus PDF fehlgeschlagen: {e}")
        return {}

def _num(x: Any) -> Optional[float]:
    """Lokale _num Funktion für Parser, da wir hier keine negativen Werte wollen."""
    if x is None: return None
    if isinstance(x, (int, float)): return abs(float(x))
    s = str(x).strip().replace("'", "").replace(" ", "").replace("’", "").replace("\xa0", "").replace(",", ".")
    s = s.replace("–", "-").replace("—", "-")
    try:
        return abs(float(s))
    except (ValueError, TypeError):
        return None