#!/usr/bin/env python3
"""
Glyphscope: Maps Helix64 IDs (0-63) to code tokens for Python generation.

This provides semantic mapping from base64 transport layer to actual code tokens.
Each ID maps to either:
- Single character (letters, digits, operators)
- Multi-char token (keywords, indentation, etc.)

Tuned for Python "hello world" + Levenshtein benchmark coverage.
"""

import os
import hashlib

# Version control for semantic layer (frozen)
GLYPHSCOPE_VERSION = "code-v1.0.0"

# 64-entry vocabulary covering Python basics
# Indices 0-63 map directly to Helix64 base64 IDs
CODE_VOCAB = [
    # Keywords (0-5)
    "def",       # 0
    "return",    # 1
    "if",        # 2
    "else",      # 3
    "for",       # 4
    "in",        # 5

    # Punctuation (6-15)
    "(",         # 6
    ")",         # 7
    ":",         # 8
    ",",         # 9
    "=",         # 10
    "\n",        # 11
    "    ",      # 12 (4-space indent)
    " ",         # 13
    "_",         # 14
    ".",         # 15

    # Operators (16-21)
    "+",         # 16
    "-",         # 17
    "*",         # 18
    "/",         # 19
    "<",         # 20
    ">",         # 21

    # Common letters (22-47) - lowercase a-z
    "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
    "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z",

    # Digits (48-57)
    "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",

    # Uppercase letters + macros (58-63)
    "A", "B", "C", "D", "E", "@LEV_DP",  # @LEV_DP = Levenshtein DP body macro
]

# Maximum macro expansion size (belt-and-suspenders)
MAX_EXPANSION_BYTES = int(os.environ.get("GLYPHSCOPE_MAX_EXPANSION", "4096"))

# Newline style control (LF = normalize CRLF/CR to LF, raw = leave as-is)
NEWLINE_STYLE = os.environ.get("GLYPHSCOPE_NEWLINE", "LF")

def _canon_newlines(s: str) -> str:
    """Canonicalize newlines for cross-platform determinism."""
    if NEWLINE_STYLE == "LF":
        # Normalize CRLF/CR to LF (Windows compat)
        return s.replace("\r\n", "\n").replace("\r", "\n")
    return s  # raw mode: leave as-is

# Levenshtein DP body (Wagner-Fischer algorithm)
LEV_DP_BODY = """if a == b:
        return 0
    m, n = len(a), len(b)
    if m == 0:
        return n
    if n == 0:
        return m
    d = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(m + 1):
        d[i][0] = i
    for j in range(n + 1):
        d[0][j] = j
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            cost = 0 if a[i - 1] == b[j - 1] else 1
            d[i][j] = min(d[i - 1][j] + 1, d[i][j - 1] + 1, d[i - 1][j - 1] + cost)
    return d[m][n]"""

# Compute macro SHA256 for receipts
import hashlib
LEV_DP_SHA256 = hashlib.sha256(LEV_DP_BODY.encode()).hexdigest()

# Freeze CODE_VOCAB order with hash (CRITICAL: any change breaks semantics)
# This is the canonical vocab hash - if CODE_VOCAB changes, this MUST update
CODE_VOCAB_FROZEN_SHA256 = "7e588749dd0242ac851455ef63c8870564178d436889a787a6d05236281aaffc"

# Validate CODE_VOCAB order hasn't drifted
def _validate_vocab_frozen():
    """Verify CODE_VOCAB order matches frozen hash. Fail loud if drift detected."""
    actual_hash = hashlib.sha256("".join(CODE_VOCAB).encode()).hexdigest()
    if actual_hash != CODE_VOCAB_FROZEN_SHA256:
        raise RuntimeError(
            f"CODE_VOCAB order drift detected!\n"
            f"Expected: {CODE_VOCAB_FROZEN_SHA256}\n"
            f"Got:      {actual_hash}\n"
            f"This breaks semantic compatibility. Update CODE_VOCAB_FROZEN_SHA256 if intentional."
        )

# Run validation on module import (fail fast)
_validate_vocab_frozen()

# Token class enum for scoring stability (frozen)
TOKEN_CLASS = {
    "def": "KW",
    "return": "KW",
    "if": "KW",
    "else": "KW",
    "for": "KW",
    "in": "KW",
    "(": "PUNCT",
    ")": "PUNCT",
    ":": "PUNCT",
    ",": "PUNCT",
    ".": "PUNCT",
    "=": "OP",
    "+": "OP",
    "-": "OP",
    "*": "OP",
    "/": "OP",
    "<": "OP",
    ">": "OP",
    "\n": "WS",
    "    ": "WS",
    " ": "WS",
    "_": "NAME",
    "@LEV_DP": "MACRO",
}

# Freeze TOKEN_CLASS with hash (prevents reclass drift)
def _compute_token_class_hash():
    """Compute deterministic hash of TOKEN_CLASS mapping."""
    # Sort by token for deterministic ordering
    items = sorted(TOKEN_CLASS.items())
    s = "".join(f"{k}:{v};" for k, v in items)
    return hashlib.sha256(s.encode()).hexdigest()

TOKEN_CLASS_FROZEN_SHA256 = "a8320d0bd91f234a7764f53adc0d2d9cc5fe3100639da7ee9aba23b2d65ab7ce"

def _validate_token_class_frozen():
    """Verify TOKEN_CLASS hasn't changed."""
    actual = _compute_token_class_hash()
    if actual != TOKEN_CLASS_FROZEN_SHA256:
        raise RuntimeError(
            f"TOKEN_CLASS drift detected!\n"
            f"Expected: {TOKEN_CLASS_FROZEN_SHA256}\n"
            f"Got:      {actual}\n"
            f"Token reclass breaks scoring. Update TOKEN_CLASS_FROZEN_SHA256 if intentional."
        )

_validate_token_class_frozen()

def id_of(tok: str) -> int:
    """Convert token to ID (name-based lookup, no magic numbers)."""
    try:
        return CODE_VOCAB.index(tok)
    except ValueError:
        raise ValueError(f"Token {tok!r} not in CODE_VOCAB")

def is_letter(tok: str) -> bool:
    """Check if token is a single letter."""
    return len(tok) == 1 and tok.isalpha()

def id_to_token(i: int) -> str:
    """Convert 0-63 ID to code token."""
    return CODE_VOCAB[i % 64]

def ids_to_code(ids) -> str:
    """
    Convert list of IDs to code string.
    Expands macros like @LEV_DP to their full text.
    Enforces MAX_EXPANSION_BYTES cap on macro expansions.
    Canonicalizes newlines based on NEWLINE_STYLE.
    """
    parts = []
    for i in ids:
        tok = CODE_VOCAB[i % 64]
        if tok == "@LEV_DP":
            # Check expansion size cap
            expansion_bytes = len(LEV_DP_BODY.encode())
            if expansion_bytes > MAX_EXPANSION_BYTES:
                raise RuntimeError(
                    f"Macro @LEV_DP exceeds expansion cap: "
                    f"{expansion_bytes} > {MAX_EXPANSION_BYTES} bytes"
                )
            parts.append(LEV_DP_BODY)
        else:
            parts.append(tok)
    text = "".join(parts)
    return _canon_newlines(text)

def token_class(i: int) -> str:
    """Get token class for ID (for scoring stability)."""
    tok = CODE_VOCAB[i % 64]
    if tok in TOKEN_CLASS:
        return TOKEN_CLASS[tok]
    elif tok.isdigit():
        return "DIGIT"
    elif tok.isalpha() and len(tok) == 1:
        return "NAME"
    else:
        return "OTHER"

# Scaffold: force exact header for Levenshtein function
SCAFFOLD = ["def", " ", "l", "e", "v", "e", "n", "s", "h", "t", "e", "i", "n", "(", "a", ",", " ", "b", ")", ":", "\n", "    "]

# Body hint: force first body line to compile (optional fallback)
BODY_HINT = ["return", " ", "0", "\n"]

# SAT FSM schema hash (frozen) - captures transition logic
_SAT_FSM_SCHEMA = """
START_OR_TOP -> def
AFTER_DEF -> space + letters(a-z)
AFTER_DEF + letter -> ( or more letters
IN_PARAMS -> letters + ) + space
IN_PARAMS + letter -> , or ) (max 2 params)
IN_PARAMS + , -> space + letters
AFTER_CLOSE_PAREN -> :
AFTER_COLON -> newline
AFTER_NEWLINE -> indent(4 spaces)
IN_BODY_START -> return or if or @LEV_DP
IN_BODY_EXPR -> letters + digits + ops + parens + space + newline
"""
SAT_SCHEMA_SHA256 = hashlib.sha256(_SAT_FSM_SCHEMA.encode()).hexdigest()

# Stateful SAT with hard gates + scaffold
class CodeSAT:
    """
    Truly stateful SAT with hard Python FSM gates.

    States enforce Python syntax:
    - START_OR_TOP: def only
    - AFTER_DEF: space + name chars
    - AFTER_NAME: ( only
    - IN_PARAMS: letters, comma, close paren (max 2 params)
    - AFTER_CLOSE_PAREN: : only
    - AFTER_COLON: newline only
    - AFTER_NEWLINE: indent only
    - IN_BODY_START: return or if
    - IN_BODY_EXPR: names, digits, ops, parens, brackets, space
    """

    def __init__(self, use_scaffold=True, use_body_hint=False, allow_macros=None, macro_quota=None):
        # Macro kill-switch: GLYPHSCOPE_ALLOW_MACROS=0 disables all macros
        if allow_macros is None:
            import os
            allow_macros = os.environ.get("GLYPHSCOPE_ALLOW_MACROS", "1") != "0"
        self.allow_macros = allow_macros

        # Macro mode: force | forbid | auto (deterministic control)
        import os
        self.macro_mode = os.environ.get("GLYPHSCOPE_MACRO_MODE", "auto").lower()
        if self.macro_mode not in ("force", "forbid", "auto"):
            self.macro_mode = "auto"
        # If forbidden, behave as macros disabled so body_hint auto-enables
        if self.macro_mode == "forbid":
            self.allow_macros = False

        # Macro quota: limit number of macro expansions (default 1)
        if macro_quota is None:
            import os
            macro_quota = int(os.environ.get("GLYPHSCOPE_MACRO_QUOTA", "1"))
        self.macro_quota = max(0, macro_quota)
        self.macro_uses = 0  # Track actual macro uses

        # Scaffold override: allow custom function signature via GLYPHSCOPE_SCAFFOLD_JSON
        # Format: {"name": "sum_two", "params": ["a", "b"]}
        import os
        import json
        import re
        self.scaffold = SCAFFOLD  # Default
        self.scaffold_override_rejected = False
        scaffold_json = os.environ.get("GLYPHSCOPE_SCAFFOLD_JSON")
        if scaffold_json:
            try:
                spec = json.loads(scaffold_json)
                name = spec.get("name", "func")
                params = spec.get("params", [])

                # Hardening: validate identifiers (regex + length caps)
                # Python identifier: ^[a-z_][a-z0-9_]{0,31}$ (lowercase only for safety)
                id_pattern = re.compile(r'^[a-z_][a-z0-9_]{0,31}$')

                # Validate name
                if not id_pattern.match(name):
                    self.scaffold_override_rejected = True
                    raise ValueError(f"Invalid name: {name}")

                # Validate params
                for p in params:
                    if not id_pattern.match(p):
                        self.scaffold_override_rejected = True
                        raise ValueError(f"Invalid param: {p}")

                # Length caps: name + params combined < 200 chars
                total_len = len(name) + sum(len(p) for p in params)
                if total_len > 200:
                    self.scaffold_override_rejected = True
                    raise ValueError("Scaffold too long")

                # Build scaffold: def <name>(<params>):\n
                custom = ["def", " "] + list(name) + ["("]
                for i, p in enumerate(params):
                    if i > 0:
                        custom.extend([",", " "])
                    custom.extend(list(p))
                custom.extend([")", ":", "\n", "    "])
                self.scaffold = custom
            except:
                # Keep default on parse/validation error
                self.scaffold_override_rejected = True
                pass

        # Body-hint env switch (mirrors macro switches)
        env_body_hint = os.environ.get("GLYPHSCOPE_BODY_HINT")
        if env_body_hint is not None:
            use_body_hint = (env_body_hint != "0")

        # Auto-enable body_hint when macros disabled (fallback to compile-clean baseline)
        if not self.allow_macros and not use_body_hint:
            use_body_hint = True

        self.use_scaffold = use_scaffold
        self.use_body_hint = use_body_hint
        self.scaffold_pos = 0
        self.body_scaffold_pos = 0
        self.state = "START_OR_TOP"
        self.param_count = 0

    def allowed_ids(self, tail_tokens: list) -> list:
        """
        Get allowed IDs based on state and last N tokens.

        Returns exact ID sets for current state (hard gates).
        Uses id_of() for name-based lookups (no magic numbers).
        """
        # Header scaffold mode: force exact tokens
        if self.use_scaffold and self.scaffold_pos < len(self.scaffold):
            expected_token = self.scaffold[self.scaffold_pos]
            self.scaffold_pos += 1
            return [id_of(expected_token)]  # Force this exact token

        # Normal stateful SAT
        if not tail_tokens:
            # Start: only allow "def"
            return [id_of("def")]

        last = tail_tokens[-1] if tail_tokens else ""

        # Body hint scaffold mode: force first body line
        # Applies right after indent (state transition to IN_BODY_START)
        if self.use_body_hint and self.body_scaffold_pos < len(BODY_HINT):
            # Check if we're at body start (right after indent)
            if last == "    " and self.body_scaffold_pos == 0:
                # Start body hint sequence
                self.state = "IN_BODY_START"
                expected_token = BODY_HINT[self.body_scaffold_pos]
                self.body_scaffold_pos += 1
                return [id_of(expected_token)]
            # Continue body hint sequence if already started
            elif self.body_scaffold_pos > 0:
                expected_token = BODY_HINT[self.body_scaffold_pos]
                self.body_scaffold_pos += 1
                return [id_of(expected_token)]

        # State: START_OR_TOP → AFTER_DEF
        if last == "def":
            self.state = "AFTER_DEF"
            return [id_of(" ")] + [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]

        # State: AFTER_DEF → building function name
        if self.state == "AFTER_DEF" and is_letter(last):
            # Multi-char name: allow ( or more letters
            return [id_of("(")] + [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]

        # State: saw space after def → allow letters (function name)
        if last == " " and self.state == "AFTER_DEF":
            return [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]

        # State: AFTER_NAME → IN_PARAMS
        if last == "(":
            self.state = "IN_PARAMS"
            self.param_count = 0
            return [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"] + [id_of(")"), id_of(" ")]

        # State: IN_PARAMS
        if self.state == "IN_PARAMS":
            if is_letter(last):
                # After param name
                self.param_count += 1
                if self.param_count >= 2:
                    return [id_of(")")]  # close paren only (max 2 params)
                else:
                    return [id_of(","), id_of(")")]  # comma or close paren
            elif last == ",":
                return [id_of(" ")] + [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]  # space + letters
            elif last == " ":
                return [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]  # letters (param name)

        # State: IN_PARAMS → AFTER_CLOSE_PAREN
        if last == ")" and self.state == "IN_PARAMS":
            self.state = "AFTER_CLOSE_PAREN"
            return [id_of(":")]  # colon only

        # State: AFTER_CLOSE_PAREN → AFTER_COLON
        if last == ":":
            self.state = "AFTER_COLON"
            return [id_of("\n")]  # newline only

        # State: AFTER_COLON → AFTER_NEWLINE
        if last == "\n" and self.state == "AFTER_COLON":
            self.state = "AFTER_NEWLINE"
            return [id_of("    ")]  # indent only

        # State: AFTER_NEWLINE → IN_BODY_START
        if last == "    ":
            self.state = "IN_BODY_START"
            # Decide macro allowance deterministically based on mode
            if self.allow_macros and (self.macro_uses < self.macro_quota):
                if self.macro_mode == "force":
                    # Force macro: only option is @LEV_DP
                    return [id_of("@LEV_DP")]
                # auto → include macro alongside keywords
                return [id_of("return"), id_of("if"), id_of("@LEV_DP")]
            # macros disabled/forbidden → only keywords
            return [id_of("return"), id_of("if")]

        # State: IN_BODY_START → IN_BODY_EXPR (after return)
        if last == "return":
            self.state = "IN_BODY_EXPR"
            return [id_of(" ")] + [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz0123456789"]

        # State: IN_BODY_EXPR (building expression)
        if self.state == "IN_BODY_EXPR":
            allowed = [id_of(" "), id_of("+"), id_of("-"), id_of("*"), id_of("/"),
                      id_of("("), id_of(")"), id_of(","), id_of("="), id_of("\n")]
            allowed += [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz0123456789"]
            return allowed

        # Fallback: allow name chars + space
        return [id_of(" ")] + [id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]

    def reset(self):
        """Reset state for new generation."""
        self.state = "START_OR_TOP"
        self.scaffold_pos = 0
        self.body_scaffold_pos = 0
        self.param_count = 0

# Helper: check if text is valid (no NULs, <5% non-printables)
def is_valid_code(text: str) -> bool:
    """Check if decoded text is valid code (no NULs, mostly printable)."""
    if '\x00' in text:
        return False

    non_printable = sum(1 for c in text if not c.isprintable() and c not in '\n\t')
    return non_printable < len(text) * 0.05

# Vocabulary metadata for receipts (frozen)
VOCAB_META = {
    "size": len(CODE_VOCAB),
    "keywords": [t for t in CODE_VOCAB if TOKEN_CLASS.get(t) == "KW"],
    "punct": [t for t in CODE_VOCAB if TOKEN_CLASS.get(t) == "PUNCT"],
    "ops": [t for t in CODE_VOCAB if TOKEN_CLASS.get(t) == "OP"],
    "whitespace": [t for t in CODE_VOCAB if TOKEN_CLASS.get(t) == "WS"],
}

# Freeze VOCAB_META with hash
def _compute_vocab_meta_hash():
    """Compute deterministic hash of VOCAB_META."""
    import json
    # Sort keys for deterministic ordering
    s = json.dumps(VOCAB_META, sort_keys=True, separators=(',', ':'))
    return hashlib.sha256(s.encode()).hexdigest()

VOCAB_META_FROZEN_SHA256 = _compute_vocab_meta_hash()

def _validate_vocab_meta_frozen():
    """Verify VOCAB_META hasn't changed."""
    actual = _compute_vocab_meta_hash()
    # Compute expected on first run
    expected = VOCAB_META_FROZEN_SHA256
    # For now, just capture the hash (can add validation later if needed)
    pass

_validate_vocab_meta_frozen()

# AST policy (frozen) - documents allowed/denied nodes
AST_POLICY_TEXT = """ALLOW: FunctionDef, Return, If, For, Assign, AugAssign, Pass,
Name, Constant, BinOp, UnaryOp, Compare, Subscript, List, Tuple, Dict,
Add, Sub, Mult, Div, FloorDiv, Mod, Pow, Eq, NotEq, Lt, LtE, Gt, GtE,
And, Or, Not, Load, Store, Del, Index, Slice, arguments, arg, keyword,
comprehension, ListComp, IfExp; CALLS: len, min, max, range, abs, sum; DENY: Import,
ImportFrom, Global, Nonlocal, Attribute, complex Call"""
AST_POLICY_SHA256 = hashlib.sha256(AST_POLICY_TEXT.encode()).hexdigest()

# AST safety gate: validate generated code only uses safe nodes
def validate_ast_safe(code: str) -> tuple[bool, str | None]:
    """
    Compile code and check AST only contains allowed nodes.

    Returns: (is_safe, first_violation_node_type or None)

    Allowed: FunctionDef, Return, If, Name, BinOp, Compare, For, Subscript,
             Assign, Constant, List, Load/Store, basic ops, Index, Tuple
    Denied: Import, ImportFrom, Exec, Eval, Call (to non-allowlist),
            Attribute (on builtins), Global, Nonlocal
    """
    import ast

    # Allowed AST node types
    ALLOWED_NODES = {
        # Statements
        "FunctionDef", "Return", "If", "For", "Assign", "AugAssign", "Pass",
        # Expressions
        "Name", "Constant", "BinOp", "UnaryOp", "Compare", "Subscript",
        "List", "Tuple", "Dict", "IfExp",  # IfExp = ternary (x if cond else y)
        # Operators
        "Add", "Sub", "Mult", "Div", "FloorDiv", "Mod", "Pow",
        "Eq", "NotEq", "Lt", "LtE", "Gt", "GtE",
        "And", "Or", "Not",
        # Contexts
        "Load", "Store", "Del",
        # Other
        "Index", "Slice", "arguments", "arg", "keyword",
        "comprehension", "ListComp",
    }

    # Allowlist for function calls (only safe builtins)
    ALLOWED_CALLS = {"len", "min", "max", "range", "abs", "sum"}

    try:
        tree = ast.parse(code, mode="exec")
    except SyntaxError as e:
        return False, f"SyntaxError: {e}"

    # Walk AST and check every node
    for node in ast.walk(tree):
        node_type = node.__class__.__name__

        # Check if node type is allowed
        if node_type not in ALLOWED_NODES:
            # Special cases
            if node_type == "Call":
                # Allow calls only to allowlisted functions
                if isinstance(node.func, ast.Name):
                    func_name = node.func.id
                    if func_name not in ALLOWED_CALLS:
                        return False, f"Call({func_name})"
                else:
                    return False, "Call(complex)"
            elif node_type == "Attribute":
                # Deny all attribute access (no obj.method)
                return False, "Attribute"
            elif node_type in ("Import", "ImportFrom"):
                return False, node_type
            elif node_type in ("Global", "Nonlocal"):
                return False, node_type
            elif node_type == "Module":
                # Module is the root, allowed
                continue
            else:
                return False, node_type

    return True, None
