#!/usr/bin/env python3
"""
Hutter Prize Compression Shim

Python shim for file I/O and benchmarking of the Lean Hutter Maximum Compression module.

Per AGENTS.md rules:
- This shim only handles file I/O, JSON serialization, and benchmarking
- Core compression logic lives in Lean (Semantics.HutterMaximumCompression)
- No compression decisions or cost calculations in Python

Usage:
    python 5-Applications/scripts/hutter_compression_shim.py input.txt output.bin
"""

import sys
import json
import time
import subprocess
from pathlib import Path
from typing import Optional


# ─── BindServer binary discovery ─────────────────────────────────────────────
# The Lean bindserver exe is declared in lakefile.toml as:
#   [[lean_exe]] name = "bindserver" root = "BindServer"
# Built with: cd 0-Core-Formalism/lean/Semantics && lake build bindserver
# Binary lands at: .lake/build/bin/bindserver

_REPO_ROOT = Path(__file__).resolve().parents[3]
_SEMANTICS_DIR = _REPO_ROOT / "0-Core-Formalism" / "lean" / "Semantics"


def _find_bindserver_binary() -> Optional[Path]:
    """Locate the compiled Lean bindserver binary.

    Search order:
      1. Canonical Semantics tree (.lake/build/bin/bindserver)
      2. Legacy tools/lean/Semantics path (used by bind_engine.py, server.js)
      3. CWD-relative fallback
    """
    candidates = [
        _SEMANTICS_DIR / ".lake" / "build" / "bin" / "bindserver",
        _REPO_ROOT / "tools" / "lean" / "Semantics" / ".lake" / "build" / "bin" / "bindserver",
        Path.cwd() / "0-Core-Formalism" / "lean" / "Semantics" / ".lake" / "build" / "bin" / "bindserver",
    ]
    for c in candidates:
        if c.exists() and c.is_file():
            return c
    return None


def _build_bindserver_request(metrics: dict) -> dict:
    """Construct a BindServer JSON-lines request for Hutter compression verification.

    Uses the "informational" metric kind (KL-divergence cost), the natural
    information-theoretic bound for compression lawfulness.

    BindServer.handleInformational checks lawfulness as JSON structural equality
    (invL == invR via json.compress). We exploit this by sending:
      left  = actual compression metrics
      right = lawful reference (emittedSymbols forced to lawfulSymbols)
    So lawful=True iff no unlawful drift (emittedSymbols == lawfulSymbols).
    The KL-divergence cost quantifies the information-theoretic distance.
    """
    left = {
        "totalPositions": float(metrics["totalPositions"]),
        "emittedSymbols": float(metrics["emittedSymbols"]),
        "lawfulSymbols": float(metrics["lawfulSymbols"]),
        "totalCost": float(metrics["totalCost"]),
        "compressionRatio": float(metrics["compressionRatio"]),
    }
    right = {
        "totalPositions": float(metrics["totalPositions"]),
        "emittedSymbols": float(metrics["lawfulSymbols"]),
        "lawfulSymbols": float(metrics["lawfulSymbols"]),
        "totalCost": float(metrics["totalCost"]),
        "compressionRatio": float(metrics["compressionRatio"]),
    }
    return {
        "metricKind": "informational",
        "left": left,
        "right": right,
        "useHistory": False,
        "historyLen": 0,
        "historyCost": 0,
        "historyTorsion": 0,
    }


def _unverified_bind_result(metrics: dict, reason: str) -> dict:
    """Return a clearly-marked unverified result when bindserver is unavailable.

    Per AGENTS.md programming-choice flow: no Python-side lawfulness computation
    is permitted. lawful=False + trace_hash="error:..." signals the gap.
    cost=0xFFFFFFFF (max Q16_16) marks the result as formally unbounded.
    """
    return {
        "left": metrics,
        "right": "hutter_compressed_output",
        "metric": {
            "cost": 0xFFFFFFFF,
            "tensor": "informational",
            "torsion": 0x00000000,
            "reference": "hutter_maximum_compression",
            "history_len": 0,
        },
        "cost": 0xFFFFFFFF,
        "witness": {
            "left_invariant": metrics.get("invariant", "unknown"),
            "right_invariant": "unverified",
            "conserved": False,
            "trace_hash": f"error:{reason}",
        },
        "lawful": False,
        "error": reason,
    }


def read_input_file(filepath: str) -> bytes:
    """Read input file (enwik8, enwik9, or test data)."""
    with open(filepath, 'rb') as f:
        return f.read()


def write_output_file(filepath: str, data: bytes) -> None:
    """Write compressed output file."""
    with open(filepath, 'wb') as f:
        f.write(data)


def call_lean_bindserver(metrics: dict) -> dict:
    """
    Call Lean BindServer to get compression bind result.

    Protocol: JSON-lines over stdin/stdout with the compiled `bindserver` exe
    (defined in 0-Core-Formalism/lean/Semantics/BindServer.lean).

    Request  (BindRequest):  metricKind, left, right, useHistory, historyLen,
                              historyCost, historyTorsion
    Response (BindResponse): cost, lawful, leftInvariant, rightInvariant,
                              traceHash, metricTensor, metricTorsion,
                              metricHistoryLen
    Error:                   {"error": "..."}

    The server reads one JSON line, writes one JSON response line, then loops.
    On EOF (stdin closed) it exits cleanly, so subprocess.run with input=
    works for one-shot calls.

    Per AGENTS.md: Python owns only I/O; Lean owns the lawfulness decision.
    If bindserver is unavailable, returns an unverified result (lawful=False)
    rather than computing lawfulness in Python.
    """
    binary = _find_bindserver_binary()
    if binary is None:
        return _unverified_bind_result(
            metrics,
            "bindserver binary not found — run: cd 0-Core-Formalism/lean/Semantics && lake build bindserver",
        )

    request = _build_bindserver_request(metrics)
    request_line = json.dumps(request, separators=(",", ":")) + "\n"

    try:
        proc = subprocess.run(
            [str(binary)],
            input=request_line,
            capture_output=True,
            text=True,
            timeout=30,
        )
    except subprocess.TimeoutExpired:
        return _unverified_bind_result(metrics, "bindserver timed out (30s)")
    except Exception as e:
        return _unverified_bind_result(metrics, f"bindserver spawn failed: {e}")

    resp_line = proc.stdout.strip()
    if not resp_line:
        stderr_snippet = proc.stderr.strip()[:200] if proc.stderr else ""
        return _unverified_bind_result(
            metrics,
            f"bindserver empty stdout (stderr={stderr_snippet})",
        )

    try:
        resp = json.loads(resp_line)
    except json.JSONDecodeError as e:
        return _unverified_bind_result(metrics, f"bindserver invalid JSON: {e}")

    if "error" in resp:
        return _unverified_bind_result(metrics, f"bindserver error: {resp['error']}")

    # Map BindResponse → shim-expected bind_result format
    return {
        "left": metrics,
        "right": "hutter_compressed_output",
        "metric": {
            "cost": resp["cost"],
            "tensor": resp["metricTensor"],
            "torsion": resp["metricTorsion"],
            "reference": "hutter_maximum_compression",
            "history_len": resp["metricHistoryLen"],
        },
        "cost": resp["cost"],
        "witness": {
            "left_invariant": resp["leftInvariant"],
            "right_invariant": resp["rightInvariant"],
            "conserved": resp["lawful"],
            "trace_hash": resp["traceHash"],
        },
        "lawful": resp["lawful"],
    }


def simulate_lean_compression(input_data: bytes) -> dict:
    """
    Simulate Lean compression pipeline for benchmarking.

    This is a placeholder that demonstrates the expected interface.
    Real implementation would call the actual Lean functions via BindServer.
    
    Returns:
        Dictionary with compression metrics
    """
    total_positions = len(input_data)
    
    # Mock compression: emit symbols at square positions (demonstration)
    emitted_symbols = 0
    lawful_symbols = 0
    total_cost = 0
    
    for pos in range(total_positions):
        k = int(pos ** 0.5)
        a = pos - k*k
        is_square = (a == 0)
        
        # Simplified emission rule: emit at squares
        if is_square:
            emitted_symbols += 1
            lawful_symbols += 1
            total_cost += 0x00000100  # Minimal cost
    
    compression_ratio = emitted_symbols / total_positions if total_positions > 0 else 0
    
    metrics = {
        "totalPositions": total_positions,
        "emittedSymbols": emitted_symbols,
        "lawfulSymbols": lawful_symbols,
        "totalCost": total_cost,
        "avgCost": total_cost / emitted_symbols if emitted_symbols > 0 else 0,
        "compressionRatio": compression_ratio,
        "invariant": "lawful_hutter_compression" if compression_ratio > 0.618 else "unlawful_hutter_drift"
    }
    
    return metrics


def benchmark_compression(input_file: str, output_file: str) -> dict:
    """
    Run compression benchmark.
    
    Args:
        input_file: Path to input file (enwik8, enwik9, etc.)
        output_file: Path for compressed output
        
    Returns:
        Benchmark results dictionary
    """
    print(f"Reading input: {input_file}")
    input_data = read_input_file(input_file)
    input_size = len(input_data)
    
    print("Running Lean compression pipeline...")
    start_time = time.time()
    
    # Get compression metrics from Lean (simulated for now)
    metrics = simulate_lean_compression(input_data)
    
    # Call Lean BindServer for formal verification
    bind_result = call_lean_bindserver(metrics)
    
    compression_time = time.time() - start_time
    
    # Write compressed output (placeholder)
    output_data = json.dumps({
        "metrics": metrics,
        "bind": bind_result
    }).encode('utf-8')
    write_output_file(output_file, output_data)
    output_size = len(output_data)
    
    # Calculate benchmark results
    results = {
        "input_file": input_file,
        "output_file": output_file,
        "input_size": input_size,
        "output_size": output_size,
        "compression_ratio": metrics["compressionRatio"],
        "compression_time": compression_time,
        "throughput_mbps": (input_size / (1024 * 1024)) / compression_time if compression_time > 0 else 0,
        "metrics": metrics,
        "bind_result": bind_result,
        "lawful": bind_result["lawful"]
    }
    
    return results


def main():
    if len(sys.argv) != 3:
        print("Usage: python hutter_compression_shim.py <input_file> <output_file>")
        sys.exit(1)
    
    input_file = sys.argv[1]
    output_file = sys.argv[2]
    
    if not Path(input_file).exists():
        print(f"Error: Input file not found: {input_file}")
        sys.exit(1)
    
    print("=== Hutter Prize Maximum Compression Benchmark ===")
    print()
    
    results = benchmark_compression(input_file, output_file)
    
    print()
    print("=== Benchmark Results ===")
    print(f"Input:  {results['input_file']} ({results['input_size']:,} bytes)")
    print(f"Output: {results['output_file']} ({results['output_size']:,} bytes)")
    print(f"Compression Ratio: {results['compression_ratio']:.6f}")
    print(f"Compression Time: {results['compression_time']:.3f}s")
    print(f"Throughput: {results['throughput_mbps']:.2f} MB/s")
    print(f"Lawful: {results['lawful']}")
    print()
    print(f"Emitted Symbols: {results['metrics']['emittedSymbols']:,} / {results['metrics']['totalPositions']:,}")
    print(f"Lawful Symbols: {results['metrics']['lawfulSymbols']:,}")
    print(f"Total Cost: {results['metrics']['totalCost']:,}")
    print(f"Invariant: {results['metrics']['invariant']}")
    
    # Save results to JSON
    results_file = output_file + ".json"
    with open(results_file, 'w') as f:
        json.dump(results, f, indent=2)
    print()
    print(f"Results saved to: {results_file}")


if __name__ == "__main__":
    main()