"""MQL (Memory Query Language) parser for LociDB.

Parses a simple SQL-like syntax into an AST represented as plain dicts.

Supported statements
--------------------

RECALL <fields> FROM memory
  WHERE SIMILAR_TO(<query>, threshold=<float>)
    AND importance > <float>
    AND ENTITY_MATCH([<entities>])
    AND timestamp WITHIN LAST <int> DAYS
    AND scope = <string>
  LIMIT <int>;

MEMORIZE {
  content: "...",
  entities: [...],
  scope: "..."
} WITH AUTO_SCORE;
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class RecallQuery:
    """Parsed RECALL statement."""

    fields: List[str] = field(default_factory=lambda: ["*"])
    similar_to: Optional[str] = None
    similarity_threshold: float = 0.0
    min_importance: float = 0.0
    entity_match: List[str] = field(default_factory=list)
    time_window_days: Optional[int] = None
    scope: Optional[str] = None
    limit: int = 10


@dataclass
class MemorizeCommand:
    """Parsed MEMORIZE statement."""

    content: str = ""
    entities: List[str] = field(default_factory=list)
    scope: str = "global"
    auto_score: bool = True


class MQLParseError(Exception):
    """Raised when MQL input cannot be parsed."""


# ── Token patterns ────────────────────────────────────────────────────

_SIMILAR_TO_RE = re.compile(
    r'SIMILAR_TO\s*\(\s*"([^"]+)"\s*(?:,\s*threshold\s*=\s*([\d.]+))?\s*\)',
    re.IGNORECASE,
)
_IMPORTANCE_RE = re.compile(r"importance\s*>\s*([\d.]+)", re.IGNORECASE)
_ENTITY_MATCH_RE = re.compile(
    r"ENTITY_MATCH\s*\(\s*\[([^\]]*)\]\s*\)", re.IGNORECASE
)
_TIME_WINDOW_RE = re.compile(
    r"timestamp\s+WITHIN\s+LAST\s+(\d+)\s+DAYS", re.IGNORECASE
)
_SCOPE_RE = re.compile(r'scope\s*=\s*"([^"]+)"', re.IGNORECASE)
_LIMIT_RE = re.compile(r"LIMIT\s+(\d+)", re.IGNORECASE)
_FIELDS_RE = re.compile(
    r"RECALL\s+(.+?)\s+FROM\s+memory", re.IGNORECASE
)

_MEMORIZE_RE = re.compile(
    r"MEMORIZE\s*(\{.*?\})\s*(WITH\s+AUTO_SCORE)?", re.IGNORECASE | re.DOTALL
)


def parse(mql: str) -> RecallQuery | MemorizeCommand:
    """Parse an MQL string and return the corresponding AST object."""
    stripped = mql.strip().rstrip(";").strip()

    if stripped.upper().startswith("RECALL"):
        return _parse_recall(stripped)
    elif stripped.upper().startswith("MEMORIZE"):
        return _parse_memorize(stripped)
    else:
        raise MQLParseError(f"Unknown MQL statement. Must start with RECALL or MEMORIZE.")


def _parse_recall(mql: str) -> RecallQuery:
    q = RecallQuery()

    # Fields
    m = _FIELDS_RE.search(mql)
    if m:
        raw = m.group(1).strip()
        q.fields = [f.strip() for f in raw.split(",")]

    # SIMILAR_TO
    m = _SIMILAR_TO_RE.search(mql)
    if m:
        q.similar_to = m.group(1)
        if m.group(2):
            q.similarity_threshold = float(m.group(2))

    # importance > X
    m = _IMPORTANCE_RE.search(mql)
    if m:
        q.min_importance = float(m.group(1))

    # ENTITY_MATCH
    m = _ENTITY_MATCH_RE.search(mql)
    if m:
        raw = m.group(1)
        # Parse comma-separated, possibly quoted strings
        entities = re.findall(r'"([^"]+)"', raw)
        if not entities:
            entities = [e.strip().strip("'") for e in raw.split(",") if e.strip()]
        q.entity_match = entities

    # timestamp WITHIN LAST N DAYS
    m = _TIME_WINDOW_RE.search(mql)
    if m:
        q.time_window_days = int(m.group(1))

    # scope
    m = _SCOPE_RE.search(mql)
    if m:
        q.scope = m.group(1)

    # LIMIT
    m = _LIMIT_RE.search(mql)
    if m:
        q.limit = int(m.group(1))

    return q


def _parse_memorize(mql: str) -> MemorizeCommand:
    m = _MEMORIZE_RE.search(mql)
    if not m:
        raise MQLParseError("Invalid MEMORIZE syntax.")

    body_raw = m.group(1)
    auto_score = m.group(2) is not None

    # The body is a JS-like object — normalise to valid JSON
    body_json = _normalise_json_body(body_raw)

    try:
        data: Dict[str, Any] = json.loads(body_json)
    except json.JSONDecodeError as exc:
        raise MQLParseError(f"Cannot parse MEMORIZE body as JSON: {exc}") from exc

    return MemorizeCommand(
        content=data.get("content", ""),
        entities=data.get("entities", []),
        scope=data.get("scope", "global"),
        auto_score=auto_score,
    )


def _normalise_json_body(raw: str) -> str:
    """Best-effort normalisation of JS-like object literals to valid JSON."""
    # Replace single quotes with double quotes
    s = raw.replace("'", '"')
    # Ensure keys are quoted: bare_key: → "bare_key":
    s = re.sub(r'(?<={|,)\s*(\w+)\s*:', r' "\1":', s)
    # Handle trailing commas
    s = re.sub(r",\s*}", "}", s)
    s = re.sub(r",\s*\]", "]", s)
    return s
