instant-bom/bom_assistant/classification/param_extractor.py

from __future__ import annotations

import json
import re
from pathlib import Path

from bom_assistant.session.models import ComponentCategory, NormalizedParams

_PATTERNS_PATH = Path(__file__).parent / "package_patterns.json"

# SI multipliers for passive values
_MULT: dict[str, float] = {
    "p": 1e-12, "n": 1e-9, "u": 1e-6, "µ": 1e-6,
    "m": 1e-3, "k": 1e3, "K": 1e3, "M": 1e6, "G": 1e9,
}

# European notation: 2k2, 4k7, 1R5, 3n3
_EURO_RE = re.compile(r"^(\d+)([pnuµmkKMGrR])(\d+)$")
# Standard: 100R, 10k, 4.7k, 100n, 1nF, 22uF, 0.001R
_STD_RE = re.compile(r"^(\d+\.?\d*)\s*([pnuµmkKMG]?)\s*([FfHhRrΩ]?)$")
# Frequency: 8MHz, 32.768kHz
_FREQ_RE = re.compile(r"(\d+\.?\d*)\s*(k|M|G)?Hz", re.IGNORECASE)


def extract_params(
    raw_value: str,
    footprint: str | None,
    category: ComponentCategory,
) -> NormalizedParams:
    params = NormalizedParams()
    params.package = _extract_package(footprint)

    cleaned = _preprocess(raw_value)
    # try each comma-separated token, not just the first
    tokens = [t.strip() for t in re.split(r"[,;]", cleaned) if t.strip()]
    primary = tokens[0] if tokens else ""

    if category == ComponentCategory.capacitor:
        for tok in tokens:
            _parse_capacitance(tok, params)
            if params.value is not None:
                break
        # fallback: scan entire string for embedded capacitance
        if params.value is None:
            _extract_capacitance_from_description(cleaned, params)
    elif category == ComponentCategory.resistor:
        for tok in tokens:
            _parse_resistance(tok, params)
            if params.value is not None:
                break
        if params.value is None:
            _extract_resistance_from_description(cleaned, params)
    elif category == ComponentCategory.inductor:
        for tok in tokens:
            _parse_inductance(tok, params)
            if params.value is not None:
                break
    elif category == ComponentCategory.crystal:
        _parse_frequency(cleaned, params)
    elif category in (ComponentCategory.ic, ComponentCategory.mosfet):
        params.part_number = raw_value.strip()
    elif category == ComponentCategory.sensor:
        m = re.search(r"(\d+\.?\d*)\s*([kKMΩRr]?)", primary.split()[-1] if primary else "")
        if m:
            _parse_resistance(m.group(0), params)

    # cross-category extractions from full string
    _extract_voltage(cleaned, params)
    _extract_current(cleaned, params)
    _extract_power(cleaned, params)
    _extract_tolerance(cleaned, params)

    return params


# ---------------------------------------------------------------------------
# Pre-processing
# ---------------------------------------------------------------------------

def _preprocess(s: str) -> str:
    # strip KiCad placeholder
    if s.strip() == "~":
        return ""
    # normalize unicode ohm
    s = s.replace("Ω", "R").replace("ω", "r")
    # European decimal: "3,3K" → "3.3K" (digit comma digit followed by unit)
    s = re.sub(r"(\d+),(\d+)([pnuµmkKMGrRFfHh])", lambda m: f"{m.group(1)}.{m.group(2)}{m.group(3)}", s)
    return s


# ---------------------------------------------------------------------------
# Capacitance
# ---------------------------------------------------------------------------

def _parse_capacitance(token: str, params: NormalizedParams) -> None:
    # ensure trailing F for bare prefix (e.g. "100n" → "100nF")
    m = re.match(r"^(\d+\.?\d*)\s*([pnuµ])[fF]?$", token, re.IGNORECASE)
    if m:
        num, prefix = m.group(1), m.group(2).lower()
        mult = _MULT.get(prefix, 1)
        params.value = float(num) * mult
        params.unit = "F"
        params.value_str = f"{num}{prefix}F"
        return

    # European: 4n7, 2u2
    m = _EURO_RE.match(token)
    if m and m.group(2).lower() in ("p", "n", "u", "µ"):
        prefix_char = m.group(2).lower()
        val = (int(m.group(1)) + int(m.group(3)) / 10 ** len(m.group(3))) * _MULT[prefix_char]
        params.value = val
        params.unit = "F"
        params.value_str = token
        return


# ---------------------------------------------------------------------------
# Resistance
# ---------------------------------------------------------------------------

def _parse_resistance(token: str, params: NormalizedParams) -> None:
    token = token.strip()

    # European: 2k2, 4k7, 1R5, 3n3
    m = _EURO_RE.match(token)
    if m:
        prefix_char = m.group(2)
        mult = _MULT.get(prefix_char, _MULT.get(prefix_char.lower(), 1))
        if prefix_char.lower() in ("r",):
            mult = 1
        val = (int(m.group(1)) + int(m.group(3)) / 10 ** len(m.group(3))) * mult
        params.value = val
        params.unit = "Ω"
        params.value_str = token
        return

    # Standard: 100R, 10k, 0.001R, 220k, 4.7k
    m = re.match(r"^(\d+\.?\d*)\s*([kKMmrRΩ]?)$", token)
    if m:
        num, suffix = m.group(1), m.group(2)
        mult_map = {"k": 1e3, "K": 1e3, "m": 1e-3, "M": 1e6, "r": 1, "R": 1, "Ω": 1, "": 1}
        mult = mult_map.get(suffix, 1)
        params.value = float(num) * mult
        params.unit = "Ω"
        params.value_str = token
        return


# ---------------------------------------------------------------------------
# Inductance
# ---------------------------------------------------------------------------

def _parse_inductance(token: str, params: NormalizedParams) -> None:
    m = re.match(r"^(\d+\.?\d*)\s*([pnuµm]?)[hH]?$", token)
    if m:
        num, prefix = m.group(1), m.group(2).lower()
        mult = _MULT.get(prefix, 1) if prefix else 1
        params.value = float(num) * mult
        params.unit = "H"
        params.value_str = f"{num}{prefix}H" if prefix else f"{num}H"
        return

    # European: 4u7
    m = _EURO_RE.match(token)
    if m and m.group(2).lower() in ("p", "n", "u", "µ", "m"):
        prefix_char = m.group(2).lower()
        val = (int(m.group(1)) + int(m.group(3)) / 10 ** len(m.group(3))) * _MULT[prefix_char]
        params.value = val
        params.unit = "H"
        params.value_str = token
        return


# ---------------------------------------------------------------------------
# Description fallback scanners (long strings like "CAP CER 0.1UF 100V X7R 0603")
# ---------------------------------------------------------------------------

def _extract_capacitance_from_description(s: str, params: NormalizedParams) -> None:
    m = re.search(r"\b(\d+\.?\d*)\s*([pnuµ])[fF]?\b", s, re.IGNORECASE)
    if m:
        num, prefix = m.group(1), m.group(2).lower()
        mult = _MULT.get(prefix, 1)
        params.value = float(num) * mult
        params.unit = "F"
        params.value_str = f"{num}{prefix}F"


def _extract_resistance_from_description(s: str, params: NormalizedParams) -> None:
    m = re.search(r"\b(\d+\.?\d*)\s*([kKMm]?)\s*[RrΩ]\b", s)
    if m:
        num, prefix = m.group(1), m.group(2)
        mult = {"k": 1e3, "K": 1e3, "m": 1e-3, "M": 1e6, "": 1}.get(prefix, 1)
        params.value = float(num) * mult
        params.unit = "Ω"
        params.value_str = m.group(0).strip()


# ---------------------------------------------------------------------------
# Frequency
# ---------------------------------------------------------------------------

def _parse_frequency(s: str, params: NormalizedParams) -> None:
    m = _FREQ_RE.search(s)
    if m:
        num = float(m.group(1))
        prefix = (m.group(2) or "").lower()
        mult = {"k": 1e3, "m": 1e6, "g": 1e9}.get(prefix, 1)
        params.value = num * mult
        params.unit = "Hz"
        params.value_str = m.group(0)


# ---------------------------------------------------------------------------
# Cross-category extractions
# ---------------------------------------------------------------------------

def _extract_voltage(s: str, params: NormalizedParams) -> None:
    if params.voltage_rating is not None:
        return
    m = re.search(r"(\d+\.?\d*)\s*[Vv]\b", s)
    if m:
        params.voltage_rating = float(m.group(1))


def _extract_current(s: str, params: NormalizedParams) -> None:
    if params.current_rating is not None:
        return
    m = re.search(r"(\d+\.?\d*)\s*[Aa]\b", s)
    if m:
        params.current_rating = float(m.group(1))


def _extract_power(s: str, params: NormalizedParams) -> None:
    if params.power_rating is not None:
        return
    m = re.search(r"(\d+\.?\d*)\s*W\b", s)
    if m:
        params.power_rating = float(m.group(1))


def _extract_tolerance(s: str, params: NormalizedParams) -> None:
    if params.tolerance is not None:
        return
    m = re.search(r"(\d+\.?\d*)\s*(ppm|%)", s, re.IGNORECASE)
    if m:
        params.tolerance = f"{m.group(1)}{m.group(2)}"


# ---------------------------------------------------------------------------
# Package — file-backed patterns + AI fallback
# ---------------------------------------------------------------------------

def _load_package_data() -> dict:
    with open(_PATTERNS_PATH, encoding="utf-8") as f:
        return json.load(f)


def _save_package_data(data: dict) -> None:
    with open(_PATTERNS_PATH, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


def _extract_package(footprint: str | None) -> str | None:
    if not footprint:
        return None

    data = _load_package_data()

    # 1. exact match
    if footprint in data["known_footprints"]:
        return data["known_footprints"][footprint]

    # 2. regex cascade
    for raw_pattern in data["patterns"]:
        m = re.search(raw_pattern, footprint)
        if m:
            return m.group(1)

    # 3. AI fallback
    package = _ai_identify_package(footprint)
    if package:
        data["known_footprints"][footprint] = package
        _save_package_data(data)
    return package


def _ai_identify_package(footprint: str) -> str | None:
    try:
        import anthropic
        client = anthropic.Anthropic()
        msg = client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=50,
            messages=[{
                "role": "user",
                "content": (
                    f"What is the standard electronic package name for this PCB footprint identifier: '{footprint}'?\n"
                    "Reply with ONLY the package name (e.g. SOIC-8, 0603, HTSOP-8, TO-220) "
                    "or 'unknown' if it cannot be determined. No explanation."
                ),
            }],
        )
        result = msg.content[0].text.strip()
        if result.lower() == "unknown" or not result:
            return None
        return result
    except Exception:
        return None