"""Keyword-rule baseline: the floor a trained model has to beat."""

from __future__ import annotations

import random

from zenith.engines.base import EngineError, Result, SpecError
from zenith.engines.registry import schemas, simulate, validate
from zenith.sft.answers import build_answer
from zenith.sft.values import NUMBER_RE, parse_number
from zenith.serve.pipeline import Response

# First matching rule wins.
RULES: list[tuple[tuple[str, ...], tuple[str, str]]] = [
    (("black hole", "event horizon", "schwarzschild", "photon sphere", "time slow", "time dilation"), ("gravity", "black_hole")),
    (("escape", "surface gravity", "leave it"), ("gravity", "escape")),
    (("force", "pull", "attraction"), ("gravity", "force")),
    (("doppler", "moving at", "moves away at", "approaching at", "velocity of", "travelling at", "radial velocity", "moved at"), ("cosmology", "doppler")),
    (("redshift", "z =", "how far", "galaxy", "receding", "mpc", "megaparsec", "light left", "hubble"), ("cosmology", "hubble")),
    (("wavelength", "photon", "spectrum", "frequency", " nm", "nanometre", "band"), ("cosmology", "spectrum")),
    (("blackbody", "peak", "kelvin", " k ", "colour", "color", "temperature of"), ("stellar", "blackbody")),
    (("main-sequence", "main sequence", "luminos", "lifetime", "live", "spectral class", "hr diagram", "bright"), ("stellar", "main_sequence")),
    (("compare", "ratio", "versus", " vs ", "side by side", "contrast", " from ", "both"), ("orbit", "compare")),
    (("orbit", "year", "period", " au", "astronomical unit"), ("orbit", "kepler")),
]


class RuleRouter:
    def route(self, question: str) -> dict:
        q = f" {question.lower()} "
        numbers = [parse_number(m.group()) for m in NUMBER_RE.finditer(question)]
        numbers = [n for n in numbers if n is not None]
        key = next((k for words, k in RULES if any(w in q for w in words)), None)
        if key is None:
            return {"sim": "none"} if not numbers else {"sim": "orbit", "op": "kepler"}
        spec = {"sim": key[0], "op": key[1]}
        if numbers:
            params = schemas()[key[0]][key[1]]
            target = next((p for p in params if p.required), params[0])
            spec[target.name] = numbers[0]
        return spec

    def ask(self, question: str, history=()) -> Response:
        spec = self.route(question)
        r = Response(spec=spec)
        try:
            result: Result = simulate(spec)
            if spec["sim"] in ("none", "unsupported"):
                r.result, r.answer = result.to_dict(), ""
                return r
            r.result = result.to_dict()
            r.answer = build_answer(validate(spec), result, random.Random(0), "direct")
        except (SpecError, EngineError) as e:
            r.error = str(e)
        return r