#!/usr/bin/env python3
"""
CDC Trace Expander - Detailed Execution Analysis
==================================================

Expands raw CDC traces into annotated step-by-step execution records.

Shows:
- Tile → Anchor → Token chain
- Logit scores and Z-score normalization
- SAT constraint filtering
- Token selection decisions
- Deterministic tiebreaker paths
"""

import json
import sys
from typing import List, Dict, Any, Optional
from pathlib import Path


def expand_cdc_trace(trace: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Expand a raw CDC trace into a full annotated execution trace.

    Args:
        trace: Raw trace dict from CDCDecoder.generate_from_anchor()

    Returns:
        List of expanded step records with full annotation
    """
    # Extract tile stream (per-step records)
    tile_stream = trace.get("tile_stream", [])

    if not tile_stream:
        return []

    expanded = []

    for i, step in enumerate(tile_stream):
        # Extract step data
        tile_sha = step.get("tile_sha256", "unknown")
        anchor_key = step.get("anchor_key", "unknown")
        choice = step.get("choice")
        symbol = step.get("symbol", "?")

        # SAT and filtering info
        sat_allowed = step.get("sat_allowed", 64)
        morph_pruned = step.get("morph_pruned", 0)
        selection_zone = step.get("selection_zone", "unknown")
        chosen_in_allowed = step.get("chosen_in_allowed", True)

        # Scoring info
        score_primary = step.get("score_primary")
        score_tiebreak = step.get("score_tiebreak")

        # Build expanded record
        expanded_step = {
            "step": i,
            "tile_sha256": tile_sha,
            "tile_sha_short": tile_sha[:16] if tile_sha != "unknown" else "unknown",
            "anchor_key": anchor_key,
            "anchor_short": anchor_key[:16] if anchor_key != "unknown" else "unknown",

            # Token selection
            "selected_token_id": choice,
            "symbol": symbol,

            # SAT constraints
            "sat_allowed_count": sat_allowed,
            "sat_constraint_strength": f"{sat_allowed}/64" if sat_allowed < 64 else "unconstrained",
            "morph_pruned_count": morph_pruned,
            "selection_zone": selection_zone,
            "sat_violation": not chosen_in_allowed,

            # Scoring
            "score_primary": score_primary,
            "score_tiebreak": score_tiebreak,

            # Metadata
            "deterministic": score_tiebreak is not None,
        }

        expanded.append(expanded_step)

    return expanded


def print_trace_summary(expanded_trace: List[Dict[str, Any]], verbose: bool = False):
    """
    Pretty-print a summary of the expanded trace.

    Args:
        expanded_trace: List of expanded step records
        verbose: If True, show full details; if False, show compact summary
    """
    if not expanded_trace:
        print("(empty trace)")
        return

    print("=" * 80)
    print("CDC TRACE EXECUTION SUMMARY")
    print("=" * 80)
    print(f"Total steps: {len(expanded_trace)}")
    print()

    if verbose:
        # Full detailed view
        for step in expanded_trace:
            print(f"Step {step['step']:03d}")
            print(f"  Tile:   {step['tile_sha_short']}...")
            print(f"  Anchor: {step['anchor_short']}...")
            print(f"  Token:  {step['selected_token_id']:2d} → '{step['symbol']}'")
            print(f"  SAT:    {step['sat_constraint_strength']} allowed")
            if step['morph_pruned_count'] > 0:
                print(f"  Morph:  {step['morph_pruned_count']} pruned")
            print(f"  Zone:   {step['selection_zone']}")
            if step['sat_violation']:
                print(f"  ⚠️  SAT VIOLATION")
            print(f"  Score:  primary={step['score_primary']}, tiebreak={step['score_tiebreak']}")
            print()
    else:
        # Compact view - just the token chain
        print("Token chain:")
        symbols = "".join(s['symbol'] for s in expanded_trace)
        # Print in chunks of 80 chars
        for i in range(0, len(symbols), 80):
            print(f"  {symbols[i:i+80]}")
        print()

        # SAT statistics
        total_constrained = sum(1 for s in expanded_trace if s['sat_allowed_count'] < 64)
        total_violations = sum(1 for s in expanded_trace if s['sat_violation'])
        avg_allowed = sum(s['sat_allowed_count'] for s in expanded_trace) / len(expanded_trace)

        print(f"SAT Statistics:")
        print(f"  Constrained steps: {total_constrained}/{len(expanded_trace)} ({100*total_constrained/len(expanded_trace):.1f}%)")
        print(f"  Average allowed: {avg_allowed:.1f}/64")
        print(f"  Violations: {total_violations}")
        print()


def export_to_csv(expanded_trace: List[Dict[str, Any]], output_path: str):
    """
    Export expanded trace to CSV for analysis in spreadsheets.

    Args:
        expanded_trace: List of expanded step records
        output_path: Path to output CSV file
    """
    import csv

    if not expanded_trace:
        print("(empty trace, skipping CSV export)")
        return

    # Define columns
    columns = [
        "step",
        "tile_sha_short",
        "anchor_short",
        "selected_token_id",
        "symbol",
        "sat_allowed_count",
        "morph_pruned_count",
        "selection_zone",
        "sat_violation",
        "score_primary",
        "score_tiebreak",
    ]

    with open(output_path, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=columns, extrasaction='ignore')
        writer.writeheader()
        writer.writerows(expanded_trace)

    print(f"✓ Exported {len(expanded_trace)} steps to {output_path}")


def export_to_html(expanded_trace: List[Dict[str, Any]], output_path: str,
                   generated_text: Optional[str] = None):
    """
    Export expanded trace to interactive HTML visualization.

    Args:
        expanded_trace: List of expanded step records
        output_path: Path to output HTML file
        generated_text: Optional decoded text to display
    """
    if not expanded_trace:
        print("(empty trace, skipping HTML export)")
        return

    html = """<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>CDC Trace Visualization</title>
    <style>
        body { font-family: 'Courier New', monospace; margin: 20px; background: #1e1e1e; color: #d4d4d4; }
        h1 { color: #569cd6; }
        .summary { background: #2d2d30; padding: 15px; margin: 20px 0; border-radius: 5px; }
        .output { background: #252526; padding: 15px; margin: 20px 0; border-radius: 5px;
                  border-left: 3px solid #569cd6; white-space: pre-wrap; }
        table { border-collapse: collapse; width: 100%; margin: 20px 0; }
        th { background: #569cd6; color: white; padding: 10px; text-align: left; }
        td { background: #2d2d30; padding: 8px; border-bottom: 1px solid #3e3e42; }
        .violation { background: #f48771 !important; }
        .constrained { color: #4ec9b0; }
        .unconstrained { color: #858585; }
        .symbol { font-weight: bold; color: #ce9178; }
        .tile { font-family: monospace; font-size: 0.9em; color: #9cdcfe; }
    </style>
</head>
<body>
    <h1>🧬 CDC Trace Visualization</h1>
"""

    # Summary stats
    total_steps = len(expanded_trace)
    total_constrained = sum(1 for s in expanded_trace if s['sat_allowed_count'] < 64)
    total_violations = sum(1 for s in expanded_trace if s['sat_violation'])
    avg_allowed = sum(s['sat_allowed_count'] for s in expanded_trace) / total_steps

    html += f"""
    <div class="summary">
        <strong>Execution Summary</strong><br>
        Total steps: {total_steps}<br>
        SAT constrained: {total_constrained}/{total_steps} ({100*total_constrained/total_steps:.1f}%)<br>
        Average allowed tokens: {avg_allowed:.1f}/64<br>
        SAT violations: {total_violations}
    </div>
"""

    # Generated text
    if generated_text:
        html += f"""
    <div class="output">
        <strong>Generated Output:</strong><br>
        {generated_text}
    </div>
"""

    # Trace table
    html += """
    <table>
        <thead>
            <tr>
                <th>Step</th>
                <th>Tile</th>
                <th>Anchor</th>
                <th>Token</th>
                <th>Symbol</th>
                <th>SAT Allowed</th>
                <th>Zone</th>
                <th>Score</th>
            </tr>
        </thead>
        <tbody>
"""

    for step in expanded_trace:
        violation_class = " class='violation'" if step['sat_violation'] else ""
        sat_class = "constrained" if step['sat_allowed_count'] < 64 else "unconstrained"

        html += f"""
            <tr{violation_class}>
                <td>{step['step']}</td>
                <td class="tile">{step['tile_sha_short']}...</td>
                <td class="tile">{step['anchor_short']}...</td>
                <td>{step['selected_token_id']}</td>
                <td class="symbol">{step['symbol']}</td>
                <td class="{sat_class}">{step['sat_constraint_strength']}</td>
                <td>{step['selection_zone']}</td>
                <td>{step['score_primary']}</td>
            </tr>
"""

    html += """
        </tbody>
    </table>
</body>
</html>
"""

    with open(output_path, 'w') as f:
        f.write(html)

    print(f"✓ Exported HTML visualization to {output_path}")


def main():
    """CLI entry point."""
    import argparse

    parser = argparse.ArgumentParser(description="Expand and visualize CDC traces")
    parser.add_argument("trace_json", help="Path to trace JSON file")
    parser.add_argument("--verbose", "-v", action="store_true",
                       help="Show full detailed trace")
    parser.add_argument("--csv", help="Export to CSV file")
    parser.add_argument("--html", help="Export to HTML visualization")

    args = parser.parse_args()

    # Load trace
    trace_path = Path(args.trace_json)
    if not trace_path.exists():
        print(f"Error: {trace_path} not found")
        return 1

    with open(trace_path) as f:
        trace = json.load(f)

    # Expand trace
    expanded = expand_cdc_trace(trace)

    # Print summary
    print_trace_summary(expanded, verbose=args.verbose)

    # Export if requested
    if args.csv:
        export_to_csv(expanded, args.csv)

    if args.html:
        generated_text = trace.get("generated_text")
        export_to_html(expanded, args.html, generated_text=generated_text)

    return 0


if __name__ == "__main__":
    sys.exit(main())
