src/vision.py

"""
vision.py — "search by photo". Claude Vision looks at an uploaded image and
describes the jewellery as structured search intent, which we feed into the SAME
semantic search pipeline as text queries.

We deliberately do NOT use CLIP: torch is blocked by Windows App Control on this
machine, so image embeddings aren't possible. Claude Vision -> text description
-> local ONNX embedding is the workaround, and it doubles as an explainable
"here's what I see" moment for the demo.

Uses the AICredits gateway (OpenAI-compatible), so images are passed as an
`image_url` data-URL, exactly like the OpenAI vision format.
"""

import json
import os

from dotenv import load_dotenv

load_dotenv()

BASE_URL = os.environ.get("AICREDITS_BASE_URL", "https://api.aicredits.in/v1")
# Vision-capable Claude model; override with AICREDITS_VISION_MODEL if needed.
VISION_MODEL = os.environ.get(
    "AICREDITS_VISION_MODEL",
    os.environ.get("AICREDITS_MODEL", "anthropic/claude-haiku-4.5"),
)
_client = None


def _get_client():
    global _client
    if _client is None:
        from openai import OpenAI
        key = os.environ.get("AICREDITS_API_KEY")
        if not key:
            return None
        _client = OpenAI(base_url=BASE_URL, api_key=key)
    return _client


# ---------------------------------------------------------------------------
# Vision provider CHAIN: AICredits (Claude, primary) -> Gemini (free fallback).
# Order via VISION_PROVIDERS env (default "aicredits,gemini").
# ---------------------------------------------------------------------------
import requests  # noqa: E402

GEMINI_VISION_MODEL = os.environ.get("GEMINI_VISION_MODEL", "gemini-flash-lite-latest")


def _parse_json(text):
    if not text:
        return None
    t = text.strip()
    if t.startswith("```"):
        t = t.split("```")[1]
        t = t[4:] if t.startswith("json") else t
        t = t.strip()
    try:
        return json.loads(t)
    except ValueError:
        s, e = t.find("{"), t.rfind("}")
        if s >= 0 and e > s:
            try:
                return json.loads(t[s:e + 1])
            except ValueError:
                return None
    return None


def _aicredits_vision(system, user_text, data_url, max_tokens=700):
    client = _get_client()
    if client is None:
        raise RuntimeError("no AICredits key")
    resp = client.chat.completions.create(
        model=VISION_MODEL, max_tokens=max_tokens, temperature=0.1,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": [
                {"type": "text", "text": user_text},
                {"type": "image_url", "image_url": {"url": data_url}},
            ]},
        ],
    )
    return resp.choices[0].message.content


def _gemini_vision(system, user_text, data_url, max_tokens=700):
    key = os.environ.get("GEMINI_API_KEY")
    if not key:
        raise RuntimeError("no GEMINI_API_KEY")
    mime, b64 = "image/jpeg", data_url
    if data_url.startswith("data:"):
        head, b64 = data_url.split(",", 1)
        mime = head.split(":")[1].split(";")[0]
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{GEMINI_VISION_MODEL}:generateContent"
    r = requests.post(
        url, headers={"x-goog-api-key": key, "Content-Type": "application/json"},
        json={"contents": [{"parts": [
            {"text": system + "\n\n" + user_text},
            {"inline_data": {"mime_type": mime, "data": b64}},
        ]}], "generationConfig": {"responseMimeType": "application/json",
                                  "maxOutputTokens": max_tokens}},
        timeout=60)
    if r.status_code != 200:
        raise RuntimeError(f"gemini {r.status_code}: {r.text[:150]}")
    return r.json()["candidates"][0]["content"]["parts"][0]["text"]


def vision_json(system, user_text, data_url, max_tokens=700):
    """Get a JSON dict from a vision model, trying providers in order. None on total failure."""
    order = os.environ.get("VISION_PROVIDERS", "aicredits,gemini").split(",")
    for prov in [p.strip().lower() for p in order]:
        try:
            if prov == "aicredits" and os.environ.get("AICREDITS_API_KEY"):
                d = _parse_json(_aicredits_vision(system, user_text, data_url, max_tokens))
            elif prov == "gemini" and os.environ.get("GEMINI_API_KEY"):
                d = _parse_json(_gemini_vision(system, user_text, data_url, max_tokens))
            else:
                continue
            if d:
                return d
        except Exception as e:  # noqa: BLE001
            print(f"[vision] {prov} failed, trying next: {str(e)[:120]}")
    return None


# ---------------------------------------------------------------------------
# Design analysis for the quote pipeline (ported from sketch-to-quote TASKS.design)
# ---------------------------------------------------------------------------
_SCALE_HINT = (
    "Typical real sizes, for scale: ring outer diameter 17-19 mm, ring band width "
    "1.5-3 mm; stud earring 5-10 mm across; pendant 8-16 mm; bracelet 55-65 mm; "
    "nose pin 4-7 mm. Melee diamonds are 0.9-1.6 mm across; a centre solitaire 3-6 mm."
)

_DESIGN_SYS = f"""You are a jewellery production analyst reading a product render.
Report only what is visible. {_SCALE_HINT}

Count stones across the WHOLE IMAGE (both earrings of a pair together). Count only
real stones SET INTO METAL with their own setting or prongs. Bright spots, specular
highlights, reflections, and the facet sparkle within one large stone are NOT
separate stones. If unsure, count fewer.

Do NOT estimate stone sizes in millimetres. Judge each group's stone diameter as a
FRACTION of the widest span of ONE piece. Halo melee might be 0.08 of the span; a
centre solitaire might be 0.45.

Set count_confidence honestly: "high" only when you could count every stone
individually; "medium"/"low" when dense enough that you estimated.

Reply with JSON only:
{{"category":"Ring|Earrings|Neckwear|Bracelet|Nose Pin",
 "pieces_visible":<separate wearable pieces shown, e.g. 2 for a pair>,
 "is_pair":<true if worn as a pair>,
 "stone_groups":[{{"count":<stones of this size in WHOLE IMAGE>,
                  "span_fraction":<stone diameter as fraction of one piece span, 0.02-0.6>,
                  "role":"centre|accent|pave|halo"}}],
 "metal":{{"form":"band|hoop|stud|pendant|chain|cuff",
          "thickness_fraction":<0.02-0.3>,"note":"<short build note>"}},
 "structural_defects":[<short strings>],
 "physics_ok":<true if manufacturable as drawn>,
 "count_confidence":"high|medium|low"}}"""


def analyze_design(data_url, note=""):
    """Vision pass over a design render -> stone count/size + structure (for the BOM).

    Returns the dict spec.buildDesignedSpec expects, or None on failure.
    Uses the vision provider chain (AICredits -> Gemini).
    """
    user_text = f"Context: {note}" if note else "Analyse this design."
    return vision_json(_DESIGN_SYS, user_text, data_url, max_tokens=700)


PROMPT = """You are a GIVA jewellery expert. Look at the jewellery in this image
and describe it for catalogue search. Return ONLY a JSON object:
- "description": one short human sentence of what you see (shown to the user).
- "query": search keywords — product type + motif + style + stones (e.g.
  "rose gold floral ring with small diamonds"). This gets embedded, so be
  descriptive but concise.
- "product_type": one of [ring, earrings, pendant, necklace, chain, bracelet,
  anklet, nosepin, mangalsutra, rakhi, toe ring, charm, coin] or null.
- "colour": "Gold" | "Silver" | "Rose Gold" | "Yellow Gold" | null (the visible
  metal colour).
- "motif": the main motif if any (heart, flower, evil eye, butterfly, star, ...) or null.
Judge only from the image. Do NOT guess price, material purity (solid vs plated),
or recipient."""


# Returned in the same shape understand_query() uses, so it drops into search().
def _empty(desc=""):
    return {
        "description": desc, "query": "", "product_type": None, "material": None,
        "colour": None, "shop_for": None, "min_price": None, "max_price": None,
        "sort": None, "exclude_stones": False, "motif": None,
    }


def describe_image(data_url: str) -> dict:
    """image data-URL -> structured search intent. Never raises.
    Uses the vision provider chain (AICredits -> Gemini)."""
    data = vision_json(PROMPT, "Describe this jewellery.", data_url, max_tokens=400)
    if not data:
        return _empty("(couldn't analyse image — no vision provider available)")
    out = _empty(data.get("description", ""))
    out["query"] = data.get("query") or data.get("description") or "jewellery"
    for k in ("product_type", "colour", "motif"):
        if data.get(k):
            out[k] = data[k]
    return out