"""Lightweight query planner that extracts common predicates from AimQL.

Scans the *processed* expression (after ``strip_query`` + ``query_add_default_expr``)
with simple regex patterns.  Returns a candidate set of run hashes from the
most selective matching index, or ``None`` to fall back to a full scan.

The planner does NOT handle AND/OR composition — it picks the single most
selective predicate.  The remaining predicates are still evaluated in-memory
by ``RestrictedPythonQuery.check()``.

Priority order (most selective first):
  1. experiment  (exact match)
  2. tag         (exact match)
  3. hparam      (equality or range — Tier 2)
  4. active      (binary split)
  5. archived    (binary split)
"""

from __future__ import annotations

import re
from typing import Any

from matyan_backend.storage.indexes import (
    lookup_by_active,
    lookup_by_archived,
    lookup_by_experiment,
    lookup_by_hparam_eq,
    lookup_by_hparam_range,
    lookup_by_tag,
)

from ._query import query_add_default_expr, strip_query

# ---------------------------------------------------------------------------
# Tier 1 patterns
# ---------------------------------------------------------------------------

_PAT_EXPERIMENT = re.compile(
    r"""run\.experiment\s*==\s*(?P<q>["'])(?P<name>.+?)(?P=q)""",
)

_PAT_TAG_IN = re.compile(
    r"""(?P<q>["'])(?P<name>.+?)(?P=q)\s+in\s+run\.tags""",
)

_PAT_ACTIVE = re.compile(
    r"""run\.active\s*==\s*(?P<val>True|False)""",
)

_PAT_ARCHIVED = re.compile(
    r"""run\.(?:is_)?archived\s*==\s*(?P<val>True|False)""",
)

# ---------------------------------------------------------------------------
# Tier 2 — hparam patterns
# ---------------------------------------------------------------------------

# Dot syntax:  run.hparams.lr == 0.001
_PAT_HPARAM_DOT = re.compile(
    r"""run\.hparams\.(?P<name>\w+)\s*(?P<op>==|!=|<=|>=|<|>)\s*(?P<val>.+?)(?=\s*(?:and|or|\)|$))""",
)

# Bracket syntax:  run["hparams"]["lr"] == 0.001
_PAT_HPARAM_BRACKET = re.compile(
    r"""run\["hparams"\]\["(?P<name>\w+)"\]\s*(?P<op>==|!=|<=|>=|<|>)\s*(?P<val>.+?)(?=\s*(?:and|or|\)|$))""",
)


def _parse_literal(raw: str) -> Any:  # noqa: ANN401
    """Best-effort parse of a Python literal from a regex capture."""
    s = raw.strip()
    if s in ("True", "False"):
        return s == "True"
    if (s.startswith('"') and s.endswith('"')) or (s.startswith("'") and s.endswith("'")):
        return s[1:-1]
    try:
        return int(s)
    except ValueError:
        pass
    try:
        return float(s)
    except ValueError:
        pass
    return s


def _plan_hparam(db: object, name: str, op: str, raw_val: str) -> list[str] | None:  # noqa: PLR0911
    """Dispatch to the appropriate hparam lookup based on *op*."""
    val = _parse_literal(raw_val)
    if op == "!=":
        return None
    if op == "==":
        return lookup_by_hparam_eq(db, name, val)
    if op == "<":
        return lookup_by_hparam_range(db, name, lo=None, hi=val)
    if op == "<=":
        results = lookup_by_hparam_range(db, name, lo=None, hi=val)
        exact = lookup_by_hparam_eq(db, name, val)
        return list(dict.fromkeys(results + exact))
    if op == ">":
        results = lookup_by_hparam_range(db, name, lo=val, hi=None)
        exact = lookup_by_hparam_eq(db, name, val)
        return [h for h in results if h not in set(exact)]
    if op == ">=":
        return lookup_by_hparam_range(db, name, lo=val, hi=None)
    return None


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


def _normalize_query(raw_query: str) -> str:
    stripped = strip_query(raw_query)
    return query_add_default_expr(stripped)


def plan_query(db: object, query: str) -> list[str] | None:  # noqa: PLR0911
    """Return candidate run hashes from index, or ``None`` for full scan."""
    expr = _normalize_query(query)
    if not expr:
        return None

    m = _PAT_EXPERIMENT.search(expr)
    if m:
        return lookup_by_experiment(db, m.group("name"))

    m = _PAT_TAG_IN.search(expr)
    if m:
        return lookup_by_tag(db, m.group("name"))

    for pat in (_PAT_HPARAM_DOT, _PAT_HPARAM_BRACKET):
        m = pat.search(expr)
        if m:
            result = _plan_hparam(db, m.group("name"), m.group("op"), m.group("val"))
            if result is not None:
                return result

    m = _PAT_ACTIVE.search(expr)
    if m:
        return lookup_by_active(db, m.group("val") == "True")

    m = _PAT_ARCHIVED.search(expr)
    if m:
        return lookup_by_archived(db, m.group("val") == "True")

    return None
