"""Parse VL model output into structured shopper predictions."""

from __future__ import annotations

import json
import re
from typing import Any

_PROB_PATTERNS = [
    re.compile(r"purchase_probability[\"']?\s*[:=]\s*([0-9]*\.?[0-9]+)", re.I),
    re.compile(r"PURCHASE_PROBABILITY:\s*([0-9]*\.?[0-9]+)", re.I),
    re.compile(r"\"purchase_probability\"\s*:\s*([0-9]*\.?[0-9]+)", re.I),
]

_REASON_PATTERNS = [
    re.compile(r"reason[\"']?\s*[:=]\s*[\"'](.+?)[\"']", re.I | re.S),
    re.compile(r"REASON:\s*(.+?)(?:\n\n|\Z)", re.I | re.S),
    re.compile(r"\"reason\"\s*:\s*\"(.+?)\"", re.I | re.S),
]


def _clamp_probability(value: float) -> float:
    return max(0.0, min(1.0, value))


def parse_prediction(text: str) -> dict[str, Any]:
    """Extract purchase_probability and reason from model text."""
    text = text.strip()
    prob: float | None = None
    reason: str | None = None

    try:
        payload = json.loads(text)
        if isinstance(payload, dict):
            if "purchase_probability" in payload:
                prob = float(payload["purchase_probability"])
            if "reason" in payload:
                reason = str(payload["reason"]).strip()
    except json.JSONDecodeError:
        pass

    if prob is None:
        for pattern in _PROB_PATTERNS:
            match = pattern.search(text)
            if match:
                prob = float(match.group(1))
                break

    if reason is None:
        for pattern in _REASON_PATTERNS:
            match = pattern.search(text)
            if match:
                reason = match.group(1).strip().strip('"').strip("'")
                break

    if prob is None:
        raise ValueError(f"Could not parse purchase_probability from model output: {text[:200]!r}")

    if reason is None:
        reason = text.splitlines()[-1].strip() if text else ""

    return {
        "purchase_probability": round(_clamp_probability(prob), 4),
        "reason": reason,
    }
