instant-bom/bom_assistant/suppliers/lcsc/categories.py

from __future__ import annotations

import json
from pathlib import Path

from bom_assistant.session.models import ComponentCategory

_CATEGORIES_PATH = Path(__file__).parent / "categories.json"

# Explicit subcategory name lists for categories where the Nuxt flat list mixes
# group headers (e.g. "Capacitors" id:495) with searchable leaves (id:1140+).
# None = fall back to _TOP_LEVEL_CATS lookup (all subcategories from that section).
_SUBCATEGORY_NAMES: dict[ComponentCategory, list[str] | None] = {
    ComponentCategory.capacitor: [
        "Aluminum - Polymer Capacitors",
        "Aluminum Electrolytic Capacitors",
        "Capacitor Networks, Arrays",
        "Ceramic Capacitors",
        "Electric Double Layer Capacitors (EDLC), Supercapacitors",
        "Film Capacitors",
        "Mica and PTFE Capacitors",
        "Motor Start, Motor Run Capacitors (AC)",
        "Niobium Oxide Capacitors",
        "Silicon Capacitors",
        "Tantalum - Polymer Capacitors",
        "Tantalum Capacitors",
        "Thin Film Capacitors",
        "Trimmers, Variable Capacitors",
    ],
    ComponentCategory.resistor: [
        "Chassis Mount Resistors",
        "Chip Resistor - Surface Mount",
        "Current Sense Resistors",
        "Resistor Networks, Arrays",
        "Specialized Resistors",
        "Through Hole Resistors",
        "Adjustable Power Resistor",
        "Rotary Potentiometers, Rheostats",
        "Slide Potentiometers",
        "Trimmer Potentiometers",
    ],
    ComponentCategory.inductor: [
        "Fixed Inductors",
        "Adjustable Inductors",
        "Arrays, Signal Transformers",
        "Wireless Charging Coils",
        "Ferrite Beads and Chips",
        "Common Mode Chokes",
    ],
    ComponentCategory.crystal: [
        "Crystals",
        "Oscillators",
        "Pin Configurable/Selectable Oscillators",
        "Programmable Oscillators",
        "Resonators",
        "VCOs (Voltage Controlled Oscillators)",
    ],
    ComponentCategory.diode: [
        "Diodes",
        "Bridge Rectifiers",
        "Rectifiers",
        "RF Diodes",
        "Variable Capacitance (Varicaps, Varactors)",
        "Zener",
        "Current Regulation - Diodes, Transistors",
    ],
    ComponentCategory.led: [
        "LED Addressable, Specialty",
        "LED Character and Numeric",
        "LED COBs, Engines, Modules, Strips",
        "LED Color Lighting",
        "LED Emitters - Infrared, UV, Visible",
        "LED Indication - Discrete",
        "LED White Lighting",
        "Circuit Board Indicators, Arrays, Light Bars, Bar Graphs",
    ],
    ComponentCategory.mosfet: [
        "Transistors",
        "FETs, MOSFETs",
        "IGBTs",
        "Bipolar (BJT)",
        "Power Driver Modules",
        "Special Purpose Transistors",
    ],
    # These have many subcategories — return all from the top-level section
    ComponentCategory.ic: None,
    ComponentCategory.connector: None,
    ComponentCategory.sensor: None,
    ComponentCategory.other: [],
}

# Top-level LCSC category names to scan when _SUBCATEGORY_NAMES entry is None
_TOP_LEVEL_CATS: dict[ComponentCategory, list[str]] = {
    ComponentCategory.ic: ["Integrated Circuits (ICs)", "Isolators"],
    ComponentCategory.connector: ["Connectors, Interconnects"],
    ComponentCategory.sensor: ["Sensors, Transducers"],
}


def load_categories() -> list[dict]:
    if not _CATEGORIES_PATH.exists():
        raise FileNotFoundError(
            f"LCSC category file not found: {_CATEGORIES_PATH}\n"
            "Run: python -m bom_assistant.suppliers.lcsc.fetch_categories"
        )
    with open(_CATEGORIES_PATH, encoding="utf-8") as f:
        return json.load(f)


def get_subcategories_for(internal_cat: ComponentCategory) -> list[dict]:
    """Return [{name, id}] dicts relevant to internal_cat for the LLM prompt."""
    names = _SUBCATEGORY_NAMES.get(internal_cat)

    if names is not None:
        if not names:
            return []
        name_set = set(names)
        cats = load_categories()
        result: list[dict] = []
        seen: set[int] = set()
        for cat in cats:
            for sub in cat.get("subcategories", []):
                if sub["name"] in name_set and sub["id"] not in seen:
                    seen.add(sub["id"])
                    result.append(sub)
        return result

    # Top-level category lookup
    target_tops = _TOP_LEVEL_CATS.get(internal_cat, [])
    if not target_tops:
        return []
    cats = load_categories()
    result = []
    seen = set()
    for cat in cats:
        if cat["category"] in target_tops:
            for sub in cat.get("subcategories", []):
                if sub["id"] not in seen:
                    seen.add(sub["id"])
                    result.append(sub)
    return result