"""
spec.py — the LLM ↔ Bouncer contract.

Two representations of one quantum intent:

* **Wire form (DSL string)** — what the (re-fine-tuned) Qwen2.5-3B model emits.
  This is intentionally the *exact* format the existing 7B adapter and training
  data already use::

      [ACTION: RANDOM] [QUBITS: 3]
      [ACTION: VQE]    [DISTANCE: 0.735]
      [ACTION: BELL]   [PAIRS: 2]
      [ACTION: GROVER] [ITEMS: 4]
      [ACTION: QAOA]   [NODES: 4] [DEPTH: 1]

  Keeping the wire format stable means the 3B re-fine-tune (Task 2) produces
  output this pipeline can consume unchanged.

* **Structured form (`QuantumSpec`)** — a typed dataclass the Bouncer and the
  rest of the pipeline work with.

`interpret_query()` is the **heuristic intent parser** used until the 3B LLM
slot is wired in (Task 2). It maps free text → DSL by keyword + number
extraction, so the deterministic pipeline is fully testable without the model.
When the LLM is present, `compile_with_llm()` replaces it and the downstream
code is untouched.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, asdict
from typing import Optional

# ── Supported actions ────────────────────────────────────────────────────────

ACTIONS = ("RANDOM", "VQE", "BELL", "GROVER", "QAOA")

# ── DSL wire format ──────────────────────────────────────────────────────────

# Ordered most-specific-first so "[ACTION: GROVER]" isn't shadowed.
_DSL_PATTERNS = {
    "RANDOM": re.compile(
        r"\[ACTION:\s*RANDOM\]\s*\[QUBITS:\s*(\d+)\s*\]", re.IGNORECASE
    ),
    "VQE": re.compile(
        r"\[ACTION:\s*VQE\]\s*\[DISTANCE:\s*(\d+(?:\.\d+)?)\s*\]", re.IGNORECASE
    ),
    "BELL": re.compile(
        r"\[ACTION:\s*BELL\]\s*\[PAIRS:\s*(\d+)\s*\]", re.IGNORECASE
    ),
    "GROVER": re.compile(
        r"\[ACTION:\s*GROVER\]\s*\[ITEMS:\s*(\d+)\s*\]", re.IGNORECASE
    ),
    "QAOA": re.compile(
        r"\[ACTION:\s*QAOA\]\s*\[NODES:\s*(\d+)\s*\]\s*\[DEPTH:\s*(\d+)\s*\]",
        re.IGNORECASE,
    ),
}


@dataclass
class QuantumSpec:
    """A validated quantum intent, before any physics checks run."""

    action: str
    qubits: Optional[int] = None        # RANDOM, BELL, GROVER
    pairs: Optional[int] = None         # BELL
    items: Optional[int] = None         # GROVER
    nodes: Optional[int] = None         # QAOA
    depth: Optional[int] = None         # QAOA
    distance: Optional[float] = None    # VQE (Å)

    def to_dsl(self) -> str:
        if self.action == "RANDOM":
            return f"[ACTION: RANDOM] [QUBITS: {self.qubits}]"
        if self.action == "VQE":
            return f"[ACTION: VQE] [DISTANCE: {self.distance:g}]"
        if self.action == "BELL":
            return f"[ACTION: BELL] [PAIRS: {self.pairs}]"
        if self.action == "GROVER":
            return f"[ACTION: GROVER] [ITEMS: {self.items}]"
        if self.action == "QAOA":
            return f"[ACTION: QAOA] [NODES: {self.nodes}] [DEPTH: {self.depth}]"
        raise ValueError(f"Unknown action: {self.action!r}")

    def to_dict(self) -> dict:
        return asdict(self)


# ── DSL → spec ───────────────────────────────────────────────────────────────

def parse_dsl(dsl_string: str) -> Optional[QuantumSpec]:
    """
    Parse a DSL wire string into a QuantumSpec, or None if it doesn't match.

    This is a *pure* parse: no physics bounds are checked here — that is the
    Bouncer's job. Returns None (not an error dict) so callers can decide how
    to surface a non-match.
    """
    raw = dsl_string.strip()
    for action, pattern in _DSL_PATTERNS.items():
        m = pattern.match(raw)
        if not m:
            continue
        if action == "RANDOM":
            return QuantumSpec(action="RANDOM", qubits=int(m.group(1)))
        if action == "VQE":
            return QuantumSpec(action="VQE", distance=float(m.group(1)))
        if action == "BELL":
            return QuantumSpec(action="BELL", pairs=int(m.group(1)))
        if action == "GROVER":
            return QuantumSpec(action="GROVER", items=int(m.group(1)))
        if action == "QAOA":
            return QuantumSpec(action="QAOA", nodes=int(m.group(1)),
                               depth=int(m.group(2)))
    return None


# ── Heuristic intent parser (fallback when no LLM is loaded) ─────────────────
#
# Keyword → action mapping, then extract the leading numeric argument. Defaults
# are conservative and intentionally pick the smallest meaningful circuit so
# the demo never silently balloons into a large one.

_KEYWORD = {
    "RANDOM": ("random", "superposition", "coin flip", "random number",
               "random bits", "randomness"),
    "VQE":    ("vqe", "ground state", "ground-state", "hydrogen", "h2",
               "molecular energy", "bond", "eigensolver", "energy of"),
    "BELL":   ("bell", "entangle", "entanglement", "epr", "bell pair"),
    "GROVER": ("grover", "grover's", "search", "database", "amplitude"),
    "QAOA":   ("qaoa", "maxcut", "max-cut", "combinatorial", "optimization",
               "optimisation"),
}


def _first_number(text: str) -> Optional[float]:
    """
    Extract the first standalone number from text.

    Uses negative look-around for letters so chemical-formula subscripts
    (the '2' in 'H2', the '3' in 'NH3', the '2O' in 'H2O') are NOT picked up.
    This matters for VQE queries like 'H2 at 0.735 Angstroms' — without it,
    the '2' in H2 is matched before '0.735' and the distance silently becomes 2.0.
    """
    nums = re.findall(r"(?<![A-Za-z])\d+(?:\.\d+)?(?![A-Za-z])", text)
    return float(nums[0]) if nums else None


def interpret_query(query: str) -> Optional[QuantumSpec]:
    """
    Map free text to a QuantumSpec using keyword + number heuristics.

    Returns None if no action keyword is recognised. This stands in for the
    Qwen2.5-3B LLM until Task 2 wires it in; the rest of the pipeline is
    identical either way because both produce a `QuantumSpec`.
    """
    q = query.lower()

    for action, keywords in _KEYWORD.items():
        if not any(kw in q for kw in keywords):
            continue
        n = _first_number(query)

        if action == "RANDOM":
            # Default 3 qubits if the user didn't specify a size.
            return QuantumSpec(action="RANDOM", qubits=int(n) if n else 3)
        if action == "VQE":
            # H₂ equilibrium bond length is the canonical default.
            return QuantumSpec(action="VQE", distance=float(n) if n else 0.735)
        if action == "BELL":
            return QuantumSpec(action="BELL", pairs=int(n) if n else 1)
        if action == "GROVER":
            # Round to a power of two-ish DB size; default 4 items.
            items = max(2, int(n) if n else 4)
            return QuantumSpec(action="GROVER", items=items)
        if action == "QAOA":
            nodes = int(n) if n else 4
            return QuantumSpec(action="QAOA", nodes=nodes, depth=1)

    return None