"""
AXON Runtime — Rust Transpiler
================================
Transpiles compute logic DSL → safe Rust source code.

The Rust transpiler converts AXON's deterministic ``logic_source``
into a Rust function with ``#[no_mangle] pub extern "C"`` linkage,
producing a C-ABI-compatible symbol that can be loaded via FFI.

Architecture (Paper §5 — Deterministic Muscle):
  logic_source  →  Rust source  →  rustc --crate-type cdylib  →  .dll/.so
                                                                     ↓
                                              ctypes.cdll  →  native execution

The transpiler enforces Linear Logic resource semantics: each input
is consumed exactly once, and the output is a strictly deterministic
morphism F: V → W with zero probabilistic component.

Security:
  - Only arithmetic operators (+, -, *, /) are emitted.
  - No unsafe blocks, no std::process, no I/O, no network.
  - The generated Rust is a pure function over f64 values.
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field


@dataclass(frozen=True)
class TranspileResult:
    """Result of transpiling logic_source to Rust."""

    rust_source: str
    fn_name: str
    param_names: tuple[str, ...]
    source_hash: str


class RustTranspiler:
    """Transpile AXON compute DSL → safe Rust source code.

    The generated Rust function:
      - Uses ``#[no_mangle] pub extern "C"`` for C ABI compatibility
      - Takes f64 parameters and returns f64
      - Contains only arithmetic operations (no I/O, no unsafe)
      - Is deterministic by construction (Linear Logic guarantee)
    """

    # Allowed operators — strict whitelist (no bitwise, no logic)
    _ALLOWED_OPS = frozenset({"+", "-", "*", "/"})

    # Pattern to match a let binding: let name = expr
    _LET_RE = re.compile(r"^let\s+(\w+)\s*=\s*(.+)$")

    # Pattern to match a return statement: return expr
    _RETURN_RE = re.compile(r"^return\s+(.+)$")

    def transpile(
        self,
        logic_source: str,
        fn_name: str,
        param_names: list[str],
    ) -> TranspileResult:
        """Transpile a logic_source string into Rust source code.

        Args:
            logic_source: The raw logic DSL (e.g. "let x = a + b\\nreturn x")
            fn_name: The compute block name (e.g. "CalculateTax")
            param_names: Input parameter names (e.g. ["amount", "rate"])

        Returns:
            TranspileResult with Rust source, function name, and hash.

        Raises:
            ValueError: If the logic_source contains unsupported constructs.
        """
        source_hash = hashlib.sha256(
            logic_source.encode("utf-8"),
        ).hexdigest()

        rust_fn_name = f"axon_compute_{self._sanitize(fn_name)}"

        # Build parameter list: all f64
        params_rust = ", ".join(
            f"{self._sanitize(p)}: f64" for p in param_names
        )

        # Transpile body lines
        body_lines = self._transpile_body(logic_source, param_names)
        body = "\n".join(f"    {line}" for line in body_lines)

        rust_source = (
            f"// Auto-generated by AXON Rust Transpiler\n"
            f"// Source hash: {source_hash}\n"
            f"// Pure deterministic morphism F: V → W\n"
            f"// Linear Logic: each resource consumed exactly once\n"
            f"\n"
            f"#[no_mangle]\n"
            f"pub extern \"C\" fn {rust_fn_name}({params_rust}) -> f64 {{\n"
            f"{body}\n"
            f"}}\n"
        )

        return TranspileResult(
            rust_source=rust_source,
            fn_name=rust_fn_name,
            param_names=tuple(param_names),
            source_hash=source_hash,
        )

    def _transpile_body(
        self,
        logic_source: str,
        param_names: list[str],
    ) -> list[str]:
        """Transpile the logic DSL body into Rust statements."""
        if not logic_source:
            raise ValueError(
                "Empty logic_source: compute blocks must contain "
                "at least one 'return' statement."
            )

        lines = logic_source.strip().splitlines()
        rust_lines: list[str] = []
        locals_declared: set[str] = set()
        has_return = False

        for line in lines:
            line = line.strip()
            if not line:
                continue

            let_match = self._LET_RE.match(line)
            if let_match:
                var_name = self._sanitize(let_match.group(1))
                expr = let_match.group(2).strip()
                rust_expr = self._transpile_expr(
                    expr, param_names, locals_declared,
                )
                rust_lines.append(f"let {var_name}: f64 = {rust_expr};")
                locals_declared.add(var_name)
                continue

            return_match = self._RETURN_RE.match(line)
            if return_match:
                expr = return_match.group(1).strip()
                rust_expr = self._transpile_expr(
                    expr, param_names, locals_declared,
                )
                rust_lines.append(f"{rust_expr}")
                has_return = True
                continue

            raise ValueError(
                f"Unsupported compute DSL statement: {line!r}"
            )

        if not has_return:
            raise ValueError(
                "Compute logic must contain a 'return' statement. "
                "Implicit return of 0.0 is not allowed for deterministic safety."
            )

        return rust_lines

    def _transpile_expr(
        self,
        expr: str,
        param_names: list[str],
        locals_declared: set[str],
    ) -> str:
        """Transpile a single DSL expression to Rust.

        Validates that only allowed identifiers and operators are used.
        """
        expr = expr.strip()
        tokens = self._tokenize_expr(expr)
        rust_tokens: list[str] = []

        for token in tokens:
            token = token.strip()
            if not token:
                continue

            # Operator
            if token in self._ALLOWED_OPS:
                rust_tokens.append(token)
                continue

            # Parentheses
            if token in ("(", ")"):
                rust_tokens.append(token)
                continue

            # Numeric literal
            if self._is_numeric(token):
                # Ensure it's an f64 literal
                if "." in token:
                    rust_tokens.append(f"{token}_f64")
                else:
                    rust_tokens.append(f"{token}.0_f64")
                continue

            # Identifier (param or local variable)
            sanitized = self._sanitize(token)
            if sanitized in {self._sanitize(p) for p in param_names}:
                rust_tokens.append(sanitized)
                continue
            if sanitized in locals_declared:
                rust_tokens.append(sanitized)
                continue

            raise ValueError(
                f"Unknown identifier in compute expression: {token!r}. "
                f"Allowed: {param_names + list(locals_declared)}"
            )

        return " ".join(rust_tokens)

    @staticmethod
    def _tokenize_expr(expr: str) -> list[str]:
        """Tokenize an arithmetic expression into atoms and operators."""
        # Split on operators and parens, keeping delimiters
        return [
            t for t in re.split(r"(\s+|[+\-*/()])", expr) if t.strip()
        ]

    @staticmethod
    def _is_numeric(token: str) -> bool:
        """Check if a token is a finite numeric literal."""
        try:
            val = float(token)
            import math
            return math.isfinite(val)
        except (ValueError, TypeError):
            return False

    @staticmethod
    def _sanitize(name: str) -> str:
        """Sanitize an identifier for Rust (alphanumeric + underscore only)."""
        return re.sub(r"[^a-zA-Z0-9_]", "_", name)
