instant-bom/bom_assistant/suppliers/category_resolver.py

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from pathlib import Path

from bom_assistant.ai.client import complete
from bom_assistant.session.models import BomRow


@dataclass
class ResolvedCategory:
    name: str
    id: int


def resolve_category(
    row: BomRow,
    subcats: list[dict],
    supplier: str,
    cache_path: Path,
) -> list[ResolvedCategory]:
    """Map a BomRow to up to 3 supplier subcategories ranked by confidence. Cache hit → instant."""
    cache = _load(cache_path)
    key = _cache_key(row)
    if key in cache:
        return _load_entry(cache[key])
    if not subcats:
        return []
    results = _ai_resolve(row, subcats, supplier)
    if results:
        cache[key] = [{"name": r.name, "id": r.id} for r in results]
        _save(cache, cache_path)
    return results


def _load_entry(e) -> list[ResolvedCategory]:
    if isinstance(e, dict):
        return [ResolvedCategory(e["name"], e["id"])]
    return [ResolvedCategory(x["name"], x["id"]) for x in e]


def _cache_key(row: BomRow) -> str:
    p = row.normalized_params
    return f"{row.category.value}|{p.package or ''}|{p.value_str or row.raw_value}"


def _fmt_subcat(s: dict) -> str:
    line = f"  {s['name']}"
    desc = s.get("description", "")
    examples = s.get("examples", [])
    if desc or examples:
        line += f" — {desc}"
        if examples:
            line += f" e.g. {', '.join(examples)}"
    return line


def _ai_resolve(row: BomRow, subcats: list[dict], supplier: str) -> list[ResolvedCategory]:
    p = row.normalized_params
    lines: list[str] = [f"  type: {row.category.value}"]
    if p.value_str:
        lines.append(f"  value: {p.value_str}")
    if p.package:
        lines.append(f"  package: {p.package}")
    if p.voltage_rating is not None:
        lines.append(f"  voltage_rating: {p.voltage_rating}V")
    if p.current_rating is not None:
        lines.append(f"  current_rating: {p.current_rating}A")
    if p.part_number:
        lines.append(f"  part_number: {p.part_number}")
    if row.raw_value and row.raw_value != p.value_str:
        lines.append(f"  raw_value: {row.raw_value!r}")

    subcat_lines = "\n".join(_fmt_subcat(s) for s in subcats)
    print(lines)
    prompt = (
        f"You are an electronics sourcing expert. Rank the top 3 {supplier.upper()} subcategories for this component, best match first.\n\n"
        "Component:\n"
        + "\n".join(lines)
        + "\n\nAvailable subcategories:\n"
        + subcat_lines
        + '\n\nReply ONLY with JSON: {"subcategories": ["<best match>", "<second best>", "<third best>"]} — use exact names from the list above, ordered by confidence'
    )
    try:
        text = complete(prompt, max_tokens=150)
        m = re.search(r'"subcategories"\s*:\s*\[([^\]]+)\]', text, re.DOTALL)
        if not m:
            return []
        names = re.findall(r'"([^"]+)"', m.group(1))
        results = []
        for name in names[:3]:
            r = _lookup(name.strip(), subcats)
            if r:
                results.append(r)
        return results
    except Exception as e:
        print(f"  [resolver error] {e}")
        return []


def _lookup(chosen: str, subcats: list[dict]) -> ResolvedCategory | None:
    cl = chosen.lower()
    for s in subcats:
        if s["name"].lower() == cl:
            return ResolvedCategory(s["name"], s["id"])
    for s in subcats:
        if cl in s["name"].lower() or s["name"].lower() in cl:
            return ResolvedCategory(s["name"], s["id"])
    return None


def _load(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}


def _save(cache: dict, path: Path) -> None:
    path.write_text(json.dumps(cache, indent=2, ensure_ascii=False), encoding="utf-8")