# ai_google_vision.py
# ------------------------------------------------------------
# Offerten-/LV-Parser für PDFs mit Gemini 2.5 pro + Fallback.
#
# Bereitstellt:
#   - extract_offer_from_pdf(pdf_path, lv_posmap=None, model_name="models/gemini-2.5-pro")
#       -> { "<pos>": {"ep": float|None, "gp": float|None}, ... }
#   - extract_lv_from_pdf(pdf_path)
#
# Hinweise:
#   - Für Gemini: Umgebungvariable GOOGLE_API_KEY setzen.
#   - Fallback nutzt pypdf + Regex, falls Gemini leer/fehlerhaft ist.
#   - Wenn lv_posmap übergeben wird, werden fehlende Positionen mit
#     {"ep": None, "gp": None} aufgefüllt (keine Key-Drops).
# ------------------------------------------------------------

from __future__ import annotations

import os
import re
import json
import logging
from pathlib import Path
from typing import Any, Dict, Optional, List

# Optional: Google Generative AI
try:
    import google.generativeai as genai
except Exception:  # pragma: no cover
    genai = None  # type: ignore

# Optional: PDF-Textauszug
try:
    from pypdf import PdfReader
except Exception:  # pragma: no cover
    PdfReader = None  # type: ignore

# ------------------------------------------------------------
# Konfiguration / Konstanten
# ------------------------------------------------------------

DEFAULT_MODEL = "models/gemini-2.5-Flash"

# Zahlenerkennung (tolerant für Schweizer/DE-Format, Tausendertrennungen etc.)
NUM = r"[+-]?\d[\d' \u00A0.,]*"
# NPK-Positionscode (z.B. 413.203)
POS_CODE = r"\b\d{3}\.\d{3}\b"


# ------------------------------------------------------------
# Hilfsfunktionen
# ------------------------------------------------------------

def _configure_model(model_name: str = DEFAULT_MODEL):
    """
    Konfiguriert google-generativeai und liefert ein GenerativeModel.
    """
    if genai is None:
        raise RuntimeError(
            "google-generativeai ist nicht installiert. "
            "Installiere es mit: pip install google-generativeai"
        )
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        raise RuntimeError(
            "GOOGLE_API_KEY ist nicht gesetzt. "
            "Setze ihn z.B. in PowerShell mit: setx GOOGLE_API_KEY \"<DEIN_KEY>\""
        )
    genai.configure(api_key=api_key)
    return genai.GenerativeModel(
        model_name,
        generation_config={
            "temperature": 0.1,
            "top_p": 0.5,
            "top_k": 32,
            "response_mime_type": "application/json",
        },
    )


def _read_pdf_text(path: Path, max_pages: Optional[int] = None) -> str:
    """
    Simpler Textauszug aus dem PDF (nur Fallback).
    """
    if PdfReader is None:
        return ""
    try:
        reader = PdfReader(str(path))
        pages = reader.pages[: (max_pages or len(reader.pages))]
        return "\n".join((p.extract_text() or "") for p in pages)
    except Exception:
        return ""


def _norm_num(val: Any) -> Optional[float]:
    """
    Normalisiert String-Zahlen (DE/CH-Format) zu float.
    """
    if val is None:
        return None
    if isinstance(val, (int, float)):
        return float(val)
    s = str(val).strip()
    if not s:
        return None

    s = s.replace("\u00A0", "").replace(" ", "").replace("'", "")

    # Wenn sowohl Punkt als auch Komma vorkommen, nimm Komma als Dezimaltrenner,
    # entferne Punkte (Tausender).
    if "," in s and "." in s:
        parts = s.split(",")
        if len(parts[-1]) in (2, 3):  # z.B. 1.234,56  /  1.234,567
            s = s.replace(".", "").replace(",", ".")
        else:
            s = s.replace(",", "")
    elif "," in s:
        s = s.replace(",", ".")

    try:
        return float(s)
    except Exception:
        return None


def _regex_prices_from_text(text: str) -> Dict[str, Dict[str, Optional[float]]]:
    """
    Heuristik-Fallback: sucht in Zeilen mit NPK-Code nach EP/GP.
    Strategie:
      1) EP/GP explizit (Schlüsselworte).
      2) Sonst die letzten zwei Zahlen der Zeile als (EP, GP) interpretieren.
    """
    result: Dict[str, Dict[str, Optional[float]]] = {}
    if not text:
        return result

    # Vorbereitete Regexe
    ep_pat = re.compile(r"(?:\bEP\b|Einheitspreis)\s*[:=]?\s*(" + NUM + r")", re.IGNORECASE)
    gp_pat = re.compile(r"(?:\bGP\b|Gesamtpreis|Total)\s*[:=]?\s*(" + NUM + r")", re.IGNORECASE)
    pos_pat = re.compile(POS_CODE)

    for rawline in text.splitlines():
        line = rawline.strip()
        posm = pos_pat.search(line)
        if not posm:
            continue
        pos = posm.group(0)

        ep = None
        gp = None

        # 1) Schlüsselwortsuche
        m_ep = ep_pat.search(line)
        m_gp = gp_pat.search(line)
        if m_ep:
            ep = _norm_num(m_ep.group(1))
        if m_gp:
            gp = _norm_num(m_gp.group(1))

        # 2) Falls nicht gefunden: letzte zwei Zahlen in der Zeile
        if ep is None or gp is None:
            nums = re.findall(NUM, line)
            if len(nums) >= 2:
                ep = ep if ep is not None else _norm_num(nums[-2])
                gp = gp if gp is not None else _norm_num(nums[-1])
            elif len(nums) == 1:
                # Wenn nur eine Zahl am Zeilenende – könnte EP oder GP sein.
                # Wir setzen sie als GP (häufiger im Spaltenlayout), EP bleibt None.
                if gp is None:
                    gp = _norm_num(nums[-1])

        # Wenn weder EP noch GP – nicht aufnehmen
        if ep is None and gp is None:
            continue

        result[pos] = {"ep": ep, "gp": gp}

    return result


def _gemini_offer_map(pdf_path: Path,
                      lv_positions: List[Dict[str, Any]],
                      model_name: str = DEFAULT_MODEL) -> Dict[str, Dict[str, Optional[float]]]:
    """
    Ruft Gemini 2.5 pro auf. Erwartet JSON:
    { "<pos>": {"ep": <float|null>, "gp": <float|null>}, ... }
    Nur die Keys aus lv_positions sollen im Ergebnis vorkommen.
    """
    model = _configure_model(model_name)

    # Datei hochladen
    try:
        # offizieller Upload-Helper der SDK
        file_obj = genai.upload_file(str(pdf_path))
    except Exception as ex:  # pragma: no cover
        raise RuntimeError(f"Upload des PDFs zu Gemini fehlgeschlagen: {ex}")

    prompt = {
        "instruction": (
            "Du bekommst eine Offerte als PDF (Datei) und eine Liste von LV-Positionen "
            "(NPK-Code, Kurztext, Menge, Einheit). "
            "Lies NUR aus dem PDF die Einheitspreise (EP) und/oder Gesamtpreise (GP) "
            "pro Position und gib ein REINES JSON zurück:\n"
            "{ \"<pos>\": {\"ep\": <float|null>, \"gp\": <float|null>}, ... }\n\n"
            "Regeln:\n"
            "- Verwende GENAU die übergebenen Positions-Keys (keine zusätzlichen/fehlenden Keys).\n"
            "- Dezimalpunkt als Trenner.\n"
            "- Wenn nur EP vorhanden → gp = null; wenn nur GP vorhanden → ep = null.\n"
            "- Wenn im PDF nichts zu finden → ep = null UND gp = null.\n"
            "- Positionen stehen als z.B. '413.203' – exakte Zuordnung.\n"
        ),
        "lv_positions": [
            {
                "pos": p.get("pos"),
                "text": (p.get("short") or p.get("text") or "")[:160],
                "qty": p.get("qty"),
                "unit": p.get("unit"),
            }
            for p in lv_positions
            if p.get("pos")
        ],
    }

    # Eingabe als Datei + JSON-Prompt
    parts = [
        file_obj,
        {"mime_type": "application/json", "text": json.dumps(prompt, ensure_ascii=False)},
    ]

    resp = model.generate_content(parts)

    # Robust Text extrahieren (keine Annahme über .text!)
    txt = ""
    cands = getattr(resp, "candidates", None) or []
    if cands and getattr(cands[0], "content", None):
        parts = getattr(cands[0].content, "parts", None) or []
        txt = "".join(getattr(p, "text", "") for p in parts)
        # Falls still beendet (z.B. finish_reason != STOP) → expliziter Fehler
        fr = getattr(cands[0], "finish_reason", None)
        # In neuerem SDK ist finish_reason evtl. Enum – auf 'STOP' prüfen, wenn vorhanden
        if fr not in (None, "STOP", 0):
            # keine 'parts' → mit separater Meldung
            if not txt.strip():
                raise RuntimeError(f"Leere Antwort vom Modell (keine 'parts'). finish_reason={fr}.")
    if not txt:
        # als Fallback versuchen wir .text
        txt = getattr(resp, "text", "") or ""

    if not txt.strip():
        raise RuntimeError("Gemini lieferte keine verwertbare JSON-Antwort (leer).")

    try:
        raw = json.loads(txt)
    except Exception as ex:
        raise RuntimeError(f"Gemini-Antwort war kein valides JSON: {ex}\nRAW (gekürzt):\n{txt[:2000]}")

    out: Dict[str, Dict[str, Optional[float]]] = {}
    for pos, vals in (raw or {}).items():
        ep = _norm_num((vals or {}).get("ep"))
        gp = _norm_num((vals or {}).get("gp"))
        out[str(pos)] = {"ep": ep, "gp": gp}
    return out


# ------------------------------------------------------------
# Öffentliche API
# ------------------------------------------------------------

def extract_offer_from_pdf(pdf_path: str | Path,
                           lv_posmap: Optional[Dict[str, Dict[str, Any]]] = None,
                           model_name: str = DEFAULT_MODEL) -> Dict[str, Dict[str, Optional[float]]]:
    """
    Liest Preise aus einer Offerte (PDF).

    Parameter:
      - pdf_path: Pfad zur Offerte (PDF).
      - lv_posmap: Optional dict {pos: {qty, unit, short/text, ...}}, um exaktes Mapping zu erzwingen.
      - model_name: Gemini-Modellname (Standard: models/gemini-2.5-pro).

    Rückgabe:
      { "<pos>": {"ep": float|None, "gp": float|None}, ... }

    Verhalten:
      - Wenn lv_posmap vorhanden + Gemini verfügbar → Gemini-Mapping.
      - Fallback (immer möglich): PDF-Text + Regex-Heuristik.
      - Fehlt ein Preis im PDF, wird ep/gp = None gesetzt (nicht 0!).
    """
    p = Path(pdf_path)

    # 1) Mit LV-Positionsliste an Gemini (wenn möglich).
    if lv_posmap and genai is not None and os.environ.get("GOOGLE_API_KEY"):
        lv_list: List[Dict[str, Any]] = []
        for pos, e in lv_posmap.items():
            lv_list.append({
                "pos": pos,
                "short": e.get("short") or e.get("text_short") or e.get("text") or "",
                "qty": e.get("qty") or e.get("quantity") or e.get("menge"),
                "unit": e.get("unit") or e.get("einheit") or e.get("unit_code"),
            })
        try:
            return _gemini_offer_map(p, lv_list, model_name=model_name)
        except Exception as ex:
            logging.warning("Gemini-Extraktion fehlgeschlagen, nutze Regex-Fallback: %s", ex)

    # 2) Fallback: Regex auf PDF-Text
    text = _read_pdf_text(p)
    prices = _regex_prices_from_text(text)

    if lv_posmap:
        # Nur die LV-Keys zurückgeben und fehlende mit None auffüllen
        filtered: Dict[str, Dict[str, Optional[float]]] = {}
        for pos in lv_posmap.keys():
            filtered[pos] = prices.get(pos, {"ep": None, "gp": None})
        return filtered

    return prices


def extract_lv_from_pdf(pdf_path: str | Path) -> Dict[str, Any]:
    """
    Grober LV-Parser (Fallback), falls mal ein LV aus PDF gelesen werden soll.
    Für euch ist CRBX die Hauptquelle – diese Funktion ist „best effort“.

    Rückgabe:
      {"posmap": {pos: {"qty": float|None, "unit": str|None, "short": str}}, "meta": {...}}
    """
    p = Path(pdf_path)
    text = _read_pdf_text(p)
    posmap: Dict[str, Dict[str, Any]] = {}

    pos_pat = re.compile(POS_CODE)

    for rawline in text.splitlines():
        line = rawline.strip()
        m = pos_pat.search(line)
        if not m:
            continue
        pos = m.group(0)

        # Naive Menge/Einheitssuche in derselben Zeile
        qty = None
        unit = None
        m_qty = re.search(rf"{POS_CODE}.*?({NUM})\s*([A-Za-zÄÖÜäöüµ²³%]+)", line)
        if m_qty:
            qty = _norm_num(m_qty.group(1))
            unit = m_qty.group(2)

        short = re.sub(POS_CODE, "", line).strip()
        posmap[pos] = {"qty": qty, "unit": unit, "short": short}

    return {"posmap": posmap, "meta": {"source": p.name}}