"""
WebAssembly (WASM) Code Generator for MechanicsDSL

Generates C code suitable for compilation to WebAssembly using Emscripten.
Includes HTML wrapper for browser-based simulation.
"""

import os
from typing import Dict, List

import sympy as sp
from sympy.printing.c import ccode

from ..utils import logger
from .base import CodeGenerator


class WasmGenerator(CodeGenerator):
    """
    Generates WebAssembly-compatible C code for browser simulations.

    Features:
    - Emscripten-compatible C code
    - HTML canvas visualization
    - JavaScript interop functions

    Example:
        >>> gen = WasmGenerator(
        ...     system_name="pendulum",
        ...     coordinates=['theta'],
        ...     parameters={'g': 9.81, 'l': 1.0},
        ...     initial_conditions={'theta': 0.5, 'theta_dot': 0.0},
        ...     equations={'theta_ddot': -g/l * sin(theta)}
        ... )
        >>> gen.generate("output/")
    """

    def __init__(
        self,
        system_name: str,
        coordinates: List[str],
        parameters: Dict[str, float],
        initial_conditions: Dict[str, float],
        equations: Dict[str, sp.Expr],
    ):

        super().__init__(system_name, coordinates, parameters, initial_conditions, equations)

    @property
    def target_name(self) -> str:
        return "wasm"

    @property
    def file_extension(self) -> str:
        return ".c"

    def generate(self, output_dir: str) -> str:
        """Generate WASM project files."""
        os.makedirs(output_dir, exist_ok=True)

        logger.info(f"Generating WASM code for {self.system_name}")

        # Generate C source
        c_file = os.path.join(output_dir, f"{self.system_name}.c")
        with open(c_file, "w", encoding="utf-8") as f:
            f.write(self._generate_c_source())

        # Generate HTML wrapper
        html_file = os.path.join(output_dir, "index.html")
        with open(html_file, "w", encoding="utf-8") as f:
            f.write(self._generate_html())

        # Generate build script
        build_file = os.path.join(output_dir, "build.sh")
        with open(build_file, "w", encoding="utf-8") as f:
            f.write(self._generate_build_script())

        logger.info(f"Generated WASM files in {output_dir}")
        return c_file

    def generate_equations(self) -> str:
        """Generate equations code."""
        lines = []
        idx = 0
        for coord in self.coordinates:
            accel_key = f"{coord}_ddot"
            lines.append(f"    dydt[{idx}] = state[{idx+1}];")
            if accel_key in self.equations:
                expr = self.equations[accel_key]
                c_expr = ccode(expr)
                lines.append(f"    dydt[{idx+1}] = {c_expr};")
            else:
                lines.append(f"    dydt[{idx+1}] = 0.0;")
            idx += 2
        return "\n".join(lines)

    def _generate_c_source(self) -> str:
        """Generate Emscripten-compatible C code."""
        state_dim = len(self.coordinates) * 2

        params = "\n".join(
            f"static const double {name} = {val};" for name, val in self.parameters.items()
        )

        init_vals = []
        for coord in self.coordinates:
            init_vals.append(str(self.initial_conditions.get(coord, 0.0)))
            init_vals.append(str(self.initial_conditions.get(f"{coord}_dot", 0.0)))
        init_str = ", ".join(init_vals)

        # State unpacking
        unpack = "\n".join(
            f"    double {c} = state[{2*i}]; double {c}_dot = state[{2*i+1}];"
            for i, c in enumerate(self.coordinates)
        )

        return f"""/*
 * WebAssembly Simulation: {self.system_name}
 * Generated by MechanicsDSL
 * 
 * Compile with Emscripten:
 *   emcc {self.system_name}.c -o {self.system_name}.js -s EXPORTED_FUNCTIONS="['_init','_step','_get_state']" -s EXPORTED_RUNTIME_METHODS="['ccall','cwrap']"
 */

#include <math.h>
#include <emscripten.h>

// Physical Parameters
{params}

#define STATE_DIM {state_dim}

// Global state
static double state[STATE_DIM];
static double t = 0.0;
static double dt = 0.01;

// Compute derivatives
void compute_derivatives(const double* state, double* dydt) {{
{unpack}

{self.generate_equations()}
}}

// RK4 integration step
void rk4_step(double* state, double dt) {{
    double k1[STATE_DIM], k2[STATE_DIM], k3[STATE_DIM], k4[STATE_DIM];
    double temp[STATE_DIM];
    int i;
    
    compute_derivatives(state, k1);
    
    for (i = 0; i < STATE_DIM; i++) temp[i] = state[i] + 0.5 * dt * k1[i];
    compute_derivatives(temp, k2);
    
    for (i = 0; i < STATE_DIM; i++) temp[i] = state[i] + 0.5 * dt * k2[i];
    compute_derivatives(temp, k3);
    
    for (i = 0; i < STATE_DIM; i++) temp[i] = state[i] + dt * k3[i];
    compute_derivatives(temp, k4);
    
    for (i = 0; i < STATE_DIM; i++) {{
        state[i] += dt * (k1[i] + 2.0*k2[i] + 2.0*k3[i] + k4[i]) / 6.0;
    }}
}}

// Exported functions
EMSCRIPTEN_KEEPALIVE
void init(void) {{
    double ic[STATE_DIM] = {{ {init_str} }};
    for (int i = 0; i < STATE_DIM; i++) {{
        state[i] = ic[i];
    }}
    t = 0.0;
}}

EMSCRIPTEN_KEEPALIVE
void step(int n) {{
    for (int i = 0; i < n; i++) {{
        rk4_step(state, dt);
        t += dt;
    }}
}}

EMSCRIPTEN_KEEPALIVE
double get_state(int idx) {{
    if (idx >= 0 && idx < STATE_DIM) {{
        return state[idx];
    }}
    return 0.0;
}}

EMSCRIPTEN_KEEPALIVE
double get_time(void) {{
    return t;
}}

EMSCRIPTEN_KEEPALIVE
void set_dt(double new_dt) {{
    dt = new_dt;
}}
"""

    def _generate_html(self) -> str:
        """Generate HTML visualization wrapper."""
        return f"""<!DOCTYPE html>
<html>
<head>
    <title>{self.system_name} - WASM Simulation</title>
    <style>
        body {{
            font-family: 'Segoe UI', Arial, sans-serif;
            max-width: 800px;
            margin: 0 auto;
            padding: 20px;
            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
            color: #eee;
            min-height: 100vh;
        }}
        h1 {{
            text-align: center;
            color: #00d4ff;
        }}
        canvas {{
            display: block;
            margin: 20px auto;
            background: #0a0a1a;
            border-radius: 10px;
            box-shadow: 0 4px 20px rgba(0,212,255,0.3);
        }}
        .controls {{
            text-align: center;
            margin: 20px 0;
        }}
        button {{
            padding: 10px 25px;
            margin: 5px;
            font-size: 16px;
            cursor: pointer;
            background: linear-gradient(45deg, #00d4ff, #0066ff);
            border: none;
            border-radius: 25px;
            color: white;
            transition: transform 0.2s;
        }}
        button:hover {{
            transform: scale(1.05);
        }}
        .info {{
            background: rgba(255,255,255,0.1);
            padding: 15px;
            border-radius: 10px;
            margin: 10px 0;
        }}
    </style>
</head>
<body>
    <h1>🌀 {self.system_name}</h1>
    <h2 style="text-align:center;color:#888;">WebAssembly Simulation</h2>
    
    <canvas id="canvas" width="600" height="400"></canvas>
    
    <div class="controls">
        <button onclick="reset()">Reset</button>
        <button onclick="togglePause()">Pause/Resume</button>
    </div>
    
    <div class="info">
        <p><strong>Time:</strong> <span id="time">0.00</span> s</p>
        <p><strong>State:</strong> <span id="state">Loading...</span></p>
    </div>
    
    <script>
        let Module;
        let running = true;
        let init_fn, step_fn, get_state_fn, get_time_fn;
        
        // Load WASM module
        var script = document.createElement('script');
        script.src = '{self.system_name}.js';
        script.onload = function() {{
            Module.onRuntimeInitialized = function() {{
                init_fn = Module.cwrap('init', null, []);
                step_fn = Module.cwrap('step', null, ['number']);
                get_state_fn = Module.cwrap('get_state', 'number', ['number']);
                get_time_fn = Module.cwrap('get_time', 'number', []);
                
                init_fn();
                requestAnimationFrame(animate);
            }};
        }};
        document.head.appendChild(script);
        
        const canvas = document.getElementById('canvas');
        const ctx = canvas.getContext('2d');
        
        function animate() {{
            if (running && step_fn) {{
                step_fn(5);  // 5 steps per frame
            }}
            draw();
            requestAnimationFrame(animate);
        }}
        
        function draw() {{
            if (!get_state_fn) return;
            
            ctx.fillStyle = 'rgba(10,10,26,0.3)';
            ctx.fillRect(0, 0, canvas.width, canvas.height);
            
            const theta = get_state_fn(0);
            const L = 150;  // Visual length
            const cx = canvas.width / 2;
            const cy = 100;
            
            const x = cx + L * Math.sin(theta);
            const y = cy + L * Math.cos(theta);
            
            // Draw rod
            ctx.strokeStyle = '#00d4ff';
            ctx.lineWidth = 3;
            ctx.beginPath();
            ctx.moveTo(cx, cy);
            ctx.lineTo(x, y);
            ctx.stroke();
            
            // Draw bob
            ctx.fillStyle = '#ff6b6b';
            ctx.beginPath();
            ctx.arc(x, y, 20, 0, 2 * Math.PI);
            ctx.fill();
            
            // Draw pivot
            ctx.fillStyle = '#888';
            ctx.beginPath();
            ctx.arc(cx, cy, 8, 0, 2 * Math.PI);
            ctx.fill();
            
            // Update info
            document.getElementById('time').textContent = get_time_fn().toFixed(2);
            document.getElementById('state').textContent = 
                'θ=' + theta.toFixed(4) + ', θ̇=' + get_state_fn(1).toFixed(4);
        }}
        
        function reset() {{
            if (init_fn) init_fn();
        }}
        
        function togglePause() {{
            running = !running;
        }}
    </script>
</body>
</html>
"""

    def _generate_build_script(self) -> str:
        """Generate Emscripten build script."""
        return f"""#!/bin/bash
# Build script for {self.system_name} WASM module

# Check for Emscripten
if ! command -v emcc &> /dev/null; then
    echo "Emscripten not found. Please install it first."
    echo "See: https://emscripten.org/docs/getting_started/downloads.html"
    exit 1
fi

# Compile
emcc {self.system_name}.c -o {self.system_name}.js \\
    -s EXPORTED_FUNCTIONS="['_init','_step','_get_state','_get_time','_set_dt']" \\
    -s EXPORTED_RUNTIME_METHODS="['ccall','cwrap']" \\
    -O3

echo "Build complete!"
echo "Open index.html in a browser to run the simulation."
"""
