#!/usr/bin/env python3
"""
Device Category Router

Rules-based classifier that predicts the device category from coarse signal
features BEFORE running the expensive per-protocol scoring loop.

Why this exists:
  The protocol database is ~300 entries, heavily dominated by weather sensors.
  Without pre-filtering, any 433 MHz RAW signal gets matched to a weather sensor
  with 65-76% confidence regardless of what it actually is. The category router
  restricts the candidate set to plausible categories, dramatically reducing
  false positives.

Usage in pattern_decoder:
  router = CategoryRouter()
  prediction = router.predict(frequency, short_us, long_us, pulse_count, preamble_type)
  candidates = [p for p in all_protocols if p.category in prediction.allowed_categories]
"""

from dataclasses import dataclass, field
from typing import List, Optional


# Canonical category names — match protocol_database.py category strings
class DeviceCategory:
    WEATHER_SENSOR   = "Weather Sensor"
    GARAGE_DOOR      = "Garage Door Opener"
    REMOTE_CONTROL   = "Remote Control"
    DOORBELL         = "Doorbell"
    TPMS             = "TPMS"
    SECURITY_SENSOR  = "Security Sensor"
    FAN_CONTROLLER   = "Fan Controller"
    SMART_METER      = "Smart Meter"
    IOT_SENSOR       = "IoT Sensor"
    UNKNOWN          = "Unknown"


# ── Decoded-protocol-name routing (KEY .sub files) ─────────────────────────
# A KEY file already carries a Flipper `Protocol:` name. For device-specific
# brands that name alone pins the category. Checked in order; first hit wins.
# (keywords, category, confidence)
_PROTOCOL_BRAND_RULES = [
    (("honeywell", "2gig", "magellan", "scher", "khan", "mastercode"),
     DeviceCategory.SECURITY_SENSOR, 0.80),
    (("came", "nice", "faac", "hormann", "somfy", "keeloq", "gate", "megacode",
      "linear", "chamberlain", "liftmaster", "marantec", "doorhan", "an-motors",
      "an_motors", "aprimatic", "beninca", "bft", "security+", "secplus", "genie",
      "clemsa", "ansonic", "sommer", "novoferm", "dooya", "alutech", "elmes",
      "nero", "hcs101", "starline", "star_line"),
     DeviceCategory.GARAGE_DOOR, 0.80),
    (("nexus", "lacrosse", "acurite", "oregon", "ambient", "infactory", "auriol",
      "bresser", "thermo", "gt-wt", "gt_wt", "wt450", "tx8300", "tx_8300",
      "wendox", "kedsum"),
     DeviceCategory.WEATHER_SENSOR, 0.75),
    (("tpms", "schrader", "pmv"),
     DeviceCategory.TPMS, 0.80),
]

# Shared line-encoder silicon reused across remotes/doorbells/fans/gates. The
# name does NOT determine device type, so we route to the whole family rather
# than dishonestly narrowing to one category.
_GENERIC_ENCODERS = (
    "princeton", "ev1527", "pt2262", "pt2264", "holtek", "ht12", "ht6",
    "smc5326", "intertechno", "rcswitch", "rc-switch", "x10", "power_smart",
)


@dataclass
class CategoryPrediction:
    """Result of category routing"""
    primary_category: str           # Best guess at device category
    confidence: float               # 0.0–1.0 confidence in this guess
    allowed_categories: List[str]   # DB categories to search (includes adjacent)
    reasoning: str                  # Human-readable explanation
    use_full_db: bool = False       # True → routing failed, fall back to all protocols


class CategoryRouter:
    """
    Rules-based device category classifier.

    Inputs  → coarse signal features (fast to compute, already available)
    Output  → CategoryPrediction with allowed_categories list

    The allowed_categories list is what gets passed to the protocol-filtering
    step. It always includes the primary category plus close neighbours so we
    don't accidentally exclude a correct match if our prediction is slightly off.
    """

    def predict(
        self,
        frequency: int,
        short_pulse_us: int,
        long_pulse_us: int,
        pulse_count: int,
        preamble_type: str = "none",
    ) -> CategoryPrediction:
        """
        Predict device category from signal features.

        Args:
            frequency:      Carrier frequency in Hz
            short_pulse_us: Extracted SHORT pulse width (µs)
            long_pulse_us:  Extracted LONG pulse width (µs)
            pulse_count:    Total number of pulses in capture
            preamble_type:  Detected preamble type from PreambleDetector

        Returns:
            CategoryPrediction
        """
        freq_mhz = frequency / 1_000_000
        ratio = long_pulse_us / short_pulse_us if short_pulse_us > 0 else 0.0

        # ── Frequency-band routing ─────────────────────────────────────────
        # 300-320 MHz  (North America: TPMS, garage doors, car fobs)
        if 300 <= freq_mhz <= 322:
            return self._route_300mhz_band(freq_mhz, short_pulse_us, pulse_count)

        # 345 MHz  (Honeywell security)
        if 343 <= freq_mhz <= 347:
            return CategoryPrediction(
                primary_category=DeviceCategory.SECURITY_SENSOR,
                confidence=0.85,
                allowed_categories=[DeviceCategory.SECURITY_SENSOR,
                                    DeviceCategory.DOORBELL],
                reasoning=f"{freq_mhz:.1f} MHz → Honeywell 345 MHz security band",
            )

        # 390 MHz  (Chamberlain/LiftMaster Security+)
        if 388 <= freq_mhz <= 392:
            return CategoryPrediction(
                primary_category=DeviceCategory.GARAGE_DOOR,
                confidence=0.90,
                allowed_categories=[DeviceCategory.GARAGE_DOOR],
                reasoning=f"{freq_mhz:.1f} MHz → LiftMaster/Chamberlain 390 MHz Security+",
            )

        # 433 MHz  (Global ISM band — most Sub-GHz IoT)
        if 433 <= freq_mhz <= 434:
            return self._route_433mhz_band(
                short_pulse_us, long_pulse_us, ratio, pulse_count, preamble_type
            )

        # 868 MHz  (European ISM: Z-Wave, smart meters, security)
        if 867 <= freq_mhz <= 869:
            return self._route_868mhz_band(short_pulse_us, pulse_count)

        # 915 MHz  (North America ISM: LoRa, industrial IoT, RFID)
        if 902 <= freq_mhz <= 928:
            return CategoryPrediction(
                primary_category=DeviceCategory.IOT_SENSOR,
                confidence=0.55,
                allowed_categories=[DeviceCategory.IOT_SENSOR,
                                    DeviceCategory.WEATHER_SENSOR,
                                    DeviceCategory.SECURITY_SENSOR],
                reasoning=f"{freq_mhz:.1f} MHz → 915 MHz North America ISM",
            )

        # Unknown frequency — fall back to full DB search
        return CategoryPrediction(
            primary_category=DeviceCategory.UNKNOWN,
            confidence=0.0,
            allowed_categories=[],
            reasoning=f"{freq_mhz:.1f} MHz — no band rule matched",
            use_full_db=True,
        )

    def route_by_protocol(
        self, protocol: Optional[str], frequency: int
    ) -> CategoryPrediction:
        """
        Route a *decoded* KEY-file to a category from its Protocol name.

        KEY .sub files have no RAW pulses to time, but they DO carry a Flipper
        `Protocol:` name. Device-specific brands (CAME, Nice, KeeLoq, Security+,
        Honeywell, ...) pin one category. Generic shared encoders (Princeton,
        EV1527, Holtek, Intertechno, ...) are the same silicon in remotes,
        doorbells, fans and gates — so we route to the whole family (a broad
        allowed set) rather than guessing a single type.

        Args:
            protocol:  Decoded Flipper protocol name (e.g. "CAME", "Princeton")
            frequency: Carrier frequency in Hz

        Returns:
            CategoryPrediction
        """
        name = (protocol or "").strip().lower()
        freq_mhz = (frequency or 0) / 1_000_000

        if not name or name in ("raw", "binraw", "unknown"):
            return CategoryPrediction(
                primary_category=DeviceCategory.UNKNOWN,
                confidence=0.0,
                allowed_categories=[],
                reasoning="no decoded protocol name",
                use_full_db=True,
            )

        # Device-specific brand → single confident category
        for keywords, category, conf in _PROTOCOL_BRAND_RULES:
            if any(k in name for k in keywords):
                return CategoryPrediction(
                    primary_category=category,
                    confidence=conf,
                    allowed_categories=[category],
                    reasoning=f"decoded protocol '{protocol}' → {category}",
                )

        # Generic shared encoder → broad family (honest 'routed' signal)
        if any(k in name for k in _GENERIC_ENCODERS):
            allowed = [DeviceCategory.REMOTE_CONTROL, DeviceCategory.DOORBELL,
                       DeviceCategory.FAN_CONTROLLER, DeviceCategory.GARAGE_DOOR]
            # 315/390 bands skew toward garage/gate; 433 is the free-for-all.
            primary = (DeviceCategory.GARAGE_DOOR if 300 <= freq_mhz <= 392
                       else DeviceCategory.REMOTE_CONTROL)
            return CategoryPrediction(
                primary_category=primary,
                confidence=0.45,
                allowed_categories=allowed,
                reasoning=(f"generic encoder '{protocol}' @ {freq_mhz:.1f} MHz "
                           "→ shared-silicon family (remote/doorbell/fan/gate)"),
            )

        # Unrecognised name — don't guess, search everything.
        return CategoryPrediction(
            primary_category=DeviceCategory.UNKNOWN,
            confidence=0.0,
            allowed_categories=[],
            reasoning=f"unrecognised protocol '{protocol}'",
            use_full_db=True,
        )

    # ── Band-specific helpers ──────────────────────────────────────────────

    def _route_300mhz_band(
        self, freq_mhz: float, short_pulse_us: int, pulse_count: int
    ) -> CategoryPrediction:
        """300–320 MHz: TPMS vs garage-door vs key-fob"""

        # 313–316 MHz is the primary TPMS band (Schrader, Continental, Pacific)
        if 313 <= freq_mhz <= 317:
            if short_pulse_us < 150:
                return CategoryPrediction(
                    primary_category=DeviceCategory.TPMS,
                    confidence=0.85,
                    allowed_categories=[DeviceCategory.TPMS],
                    reasoning=f"{freq_mhz:.1f} MHz + short_pulse={short_pulse_us}µs → TPMS",
                )
            # 315 MHz can also be garage doors / key fobs
            return CategoryPrediction(
                primary_category=DeviceCategory.GARAGE_DOOR,
                confidence=0.75,
                allowed_categories=[DeviceCategory.GARAGE_DOOR,
                                    DeviceCategory.REMOTE_CONTROL,
                                    DeviceCategory.TPMS],
                reasoning=f"315 MHz + larger pulses → garage door or remote",
            )

        # 318 MHz  (Linear MegaCode)
        if 317 <= freq_mhz <= 320:
            return CategoryPrediction(
                primary_category=DeviceCategory.GARAGE_DOOR,
                confidence=0.80,
                allowed_categories=[DeviceCategory.GARAGE_DOOR],
                reasoning=f"{freq_mhz:.1f} MHz → Linear MegaCode / gate opener",
            )

        return CategoryPrediction(
            primary_category=DeviceCategory.REMOTE_CONTROL,
            confidence=0.50,
            allowed_categories=[DeviceCategory.REMOTE_CONTROL,
                                DeviceCategory.GARAGE_DOOR,
                                DeviceCategory.TPMS],
            reasoning=f"{freq_mhz:.1f} MHz 300-band, no specific rule",
        )

    def _route_433mhz_band(
        self,
        short_pulse_us: int,
        long_pulse_us: int,
        ratio: float,
        pulse_count: int,
        preamble_type: str,
    ) -> CategoryPrediction:
        """
        433 MHz is the busiest band.  We use timing + preamble + pulse count
        to differentiate the major device families.

        Key discriminators (empirically derived):
          ratio ~ 2:1 + alternating/long_burst preamble + high pulse count → weather sensor
          ratio ~ 3:1 + sync_word preamble + low pulse count               → remote/garage
          ratio ~ 1:1 (Manchester)                                          → Oregon Sci / TPMS
          very short pulses (< 200µs)                                       → Manchester weather
          medium pulse count + no/sync preamble                             → doorbell / fan
        """

        # ── Manchester-like (ratio ≈ 1:1, short pulses) ───────────────────
        # Oregon Scientific, some TPMS that land near 433 MHz
        if 0.8 <= ratio <= 1.3 and short_pulse_us < 600:
            return CategoryPrediction(
                primary_category=DeviceCategory.WEATHER_SENSOR,
                confidence=0.75,
                allowed_categories=[DeviceCategory.WEATHER_SENSOR],
                reasoning=f"433 MHz, ratio={ratio:.2f} (≈1:1) → Manchester weather sensor",
            )

        # ── EV1527 / Princeton family (ratio > 2.3) ────────────────────────
        # IMPORTANT: Weather sensors at 433 MHz use 2:1 PWM ratio (short=0, long=1).
        # Any ratio > 2.3 is NOT a standard weather sensor protocol.
        # This rule fires regardless of pulse count (repeated remotes look high-count).
        if ratio > 2.3:
            if short_pulse_us < 400:
                return CategoryPrediction(
                    primary_category=DeviceCategory.REMOTE_CONTROL,
                    confidence=0.80,
                    allowed_categories=[DeviceCategory.REMOTE_CONTROL,
                                        DeviceCategory.DOORBELL,
                                        DeviceCategory.FAN_CONTROLLER,
                                        DeviceCategory.GARAGE_DOOR],
                    reasoning=(
                        f"433 MHz, ratio={ratio:.2f} (>2.3), short={short_pulse_us}µs "
                        f"→ EV1527/Princeton family (remote/fan/doorbell)"
                    ),
                )
            # Longer pulses at high ratio → Princeton, FAAC, older garage openers
            return CategoryPrediction(
                primary_category=DeviceCategory.GARAGE_DOOR,
                confidence=0.72,
                allowed_categories=[DeviceCategory.GARAGE_DOOR,
                                    DeviceCategory.REMOTE_CONTROL,
                                    DeviceCategory.DOORBELL],
                reasoning=(
                    f"433 MHz, ratio={ratio:.2f}, short={short_pulse_us}µs (>400µs) "
                    "→ Princeton/FAAC/large garage remote"
                ),
            )

        # ── PWM weather sensors (ratio 1.8–2.3, lots of pulses, preamble) ─
        if (1.8 <= ratio <= 2.3
                and pulse_count >= 100
                and preamble_type in ("alternating", "long_burst", "none")):
            return CategoryPrediction(
                primary_category=DeviceCategory.WEATHER_SENSOR,
                confidence=0.80,
                allowed_categories=[DeviceCategory.WEATHER_SENSOR],
                reasoning=(
                    f"433 MHz, ratio={ratio:.2f}, pulses={pulse_count}, "
                    f"preamble={preamble_type} → PWM weather sensor"
                ),
            )

        # ── Garage doors with 2:1 ratio (LiftMaster, short burst) ─────────
        # Garage door presses are short (one button press = <100 pulses per repeat)
        # They also tend to have a distinctive sync gap before data
        if 1.5 <= ratio <= 2.5 and pulse_count < 100:
            return CategoryPrediction(
                primary_category=DeviceCategory.GARAGE_DOOR,
                confidence=0.65,
                allowed_categories=[DeviceCategory.GARAGE_DOOR,
                                    DeviceCategory.REMOTE_CONTROL,
                                    DeviceCategory.DOORBELL],
                reasoning=(
                    f"433 MHz, ratio={ratio:.2f}, pulses={pulse_count} (low) "
                    "→ garage/remote short burst"
                ),
            )

        # ── Doorbells (ratio 1.5–2.3, short_pulse 150–350µs, low count) ───
        if 1.5 <= ratio <= 2.3 and short_pulse_us < 350 and pulse_count < 150:
            return CategoryPrediction(
                primary_category=DeviceCategory.DOORBELL,
                confidence=0.70,
                allowed_categories=[DeviceCategory.DOORBELL,
                                    DeviceCategory.REMOTE_CONTROL,
                                    DeviceCategory.FAN_CONTROLLER],
                reasoning=(
                    f"433 MHz, ratio={ratio:.2f}, short={short_pulse_us}µs, "
                    f"pulses={pulse_count} → doorbell or short remote"
                ),
            )

        # ── Catch-all: ratio unclear — broad search
        if pulse_count >= 150:
            return CategoryPrediction(
                primary_category=DeviceCategory.WEATHER_SENSOR,
                confidence=0.55,
                allowed_categories=[DeviceCategory.WEATHER_SENSOR,
                                    DeviceCategory.SECURITY_SENSOR],
                reasoning=(
                    f"433 MHz, ratio={ratio:.2f}, pulses={pulse_count} (high) "
                    "→ likely sensor (low confidence)"
                ),
            )

        return CategoryPrediction(
            primary_category=DeviceCategory.REMOTE_CONTROL,
            confidence=0.40,
            allowed_categories=[DeviceCategory.REMOTE_CONTROL,
                                DeviceCategory.DOORBELL,
                                DeviceCategory.FAN_CONTROLLER,
                                DeviceCategory.GARAGE_DOOR,
                                DeviceCategory.WEATHER_SENSOR],
            reasoning=(
                f"433 MHz, ratio={ratio:.2f}, pulses={pulse_count} — "
                "no clear rule, broad search"
            ),
            use_full_db=False,
        )

    def _route_868mhz_band(
        self, short_pulse_us: int, pulse_count: int
    ) -> CategoryPrediction:
        """868 MHz: Z-Wave, security sensors, smart meters, LoRa"""

        # Z-Wave uses 100µs pulses at 868.42 MHz (GFSK, not OOK)
        # We can only classify as security/smart-home
        if short_pulse_us < 200:
            return CategoryPrediction(
                primary_category=DeviceCategory.SECURITY_SENSOR,
                confidence=0.70,
                allowed_categories=[DeviceCategory.SECURITY_SENSOR,
                                    DeviceCategory.SMART_METER],
                reasoning=f"868 MHz + fast pulses ({short_pulse_us}µs) → Z-Wave or security",
            )

        # European weather sensors also appear at 868 MHz
        if pulse_count >= 100:
            return CategoryPrediction(
                primary_category=DeviceCategory.WEATHER_SENSOR,
                confidence=0.65,
                allowed_categories=[DeviceCategory.WEATHER_SENSOR,
                                    DeviceCategory.SECURITY_SENSOR,
                                    DeviceCategory.SMART_METER],
                reasoning=f"868 MHz + high pulse count → EU weather sensor or meter",
            )

        return CategoryPrediction(
            primary_category=DeviceCategory.SECURITY_SENSOR,
            confidence=0.60,
            allowed_categories=[DeviceCategory.SECURITY_SENSOR,
                                DeviceCategory.SMART_METER,
                                DeviceCategory.WEATHER_SENSOR],
            reasoning=f"868 MHz general → security / smart meter",
        )


# Singleton
_router: Optional[CategoryRouter] = None


def get_category_router() -> CategoryRouter:
    global _router
    if _router is None:
        _router = CategoryRouter()
    return _router
