"""
JavaScript Code Generator for MechanicsDSL

Generates standalone JavaScript simulation code for Node.js or browsers.
"""

from sympy.printing.jscode import jscode

from ..utils import logger
from .base import CodeGenerator


class JavaScriptGenerator(CodeGenerator):
    """
    Generates JavaScript simulation code.

    Produces ES6+ modules with built-in RK4 integrator for browser use.
    """

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

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

    def generate(self, output_file: str = "simulation.js") -> str:
        """Generate JavaScript simulation code."""
        logger.info(f"Generating JavaScript code for {self.system_name}")

        code = self._generate_code()

        with open(output_file, "w") as f:
            f.write(code)

        logger.info(f"Successfully wrote {output_file}")
        return output_file

    def generate_equations(self) -> str:
        """Generate JavaScript code for equations."""
        lines = []
        idx = 0
        for coord in self.coordinates:
            accel_key = f"{coord}_ddot"
            lines.append(f"    dydt[{idx}] = y[{idx+1}];  // d{coord}/dt = {coord}_dot")
            if accel_key in self.equations:
                expr = self.equations[accel_key]
                js_expr = jscode(expr)
                lines.append(f"    dydt[{idx+1}] = {js_expr};  // d{coord}_dot/dt")
            else:
                lines.append(f"    dydt[{idx+1}] = 0.0;")
            idx += 2
        return "\n".join(lines)

    def _generate_code(self) -> str:
        """Generate complete JavaScript simulation module."""
        # Parameter definitions
        param_lines = []
        for name, val in self.parameters.items():
            param_lines.append(f"const {name} = {val};")
        param_str = "\n".join(param_lines)

        # State variable unpacking
        unpack_lines = []
        idx = 0
        for coord in self.coordinates:
            unpack_lines.append(f"    const {coord} = y[{idx}];")
            unpack_lines.append(f"    const {coord}_dot = y[{idx+1}];")
            idx += 2
        unpack_str = "\n".join(unpack_lines)

        # Equations
        eq_str = self.generate_equations()

        # Initial conditions
        init_vals = []
        for coord in self.coordinates:
            pos = self.initial_conditions.get(coord, 0.0)
            vel = self.initial_conditions.get(f"{coord}_dot", 0.0)
            init_vals.extend([str(pos), str(vel)])
        init_str = ", ".join(init_vals)

        state_dim = len(self.coordinates) * 2

        template = f"""/**
 * Auto-generated simulation: {self.system_name}
 * Generated by MechanicsDSL
 * 
 * Can run in Node.js or browsers (ES6+ modules)
 */

// Physical Parameters
{param_str}

const DIM = {state_dim};

/**
 * Equations of motion
 * @param {{number}} t - Current time
 * @param {{number[]}} y - State vector
 * @returns {{number[]}} - Derivatives
 */
function equationsOfMotion(t, y) {{
    const dydt = new Array(DIM).fill(0);
    
    // Unpack state
{unpack_str}
    
    // Compute derivatives
{eq_str}
    
    return dydt;
}}

/**
 * RK4 integration step
 */
function rk4Step(t, y, dt) {{
    const k1 = equationsOfMotion(t, y);
    const k2 = equationsOfMotion(t + 0.5*dt, y.map((yi, i) => yi + 0.5*dt*k1[i]));
    const k3 = equationsOfMotion(t + 0.5*dt, y.map((yi, i) => yi + 0.5*dt*k2[i]));
    const k4 = equationsOfMotion(t + dt, y.map((yi, i) => yi + dt*k3[i]));
    
    return y.map((yi, i) => yi + dt/6 * (k1[i] + 2*k2[i] + 2*k3[i] + k4[i]));
}}

/**
 * Run simulation
 * @param {{number}} tEnd - End time
 * @param {{number}} dt - Time step
 * @returns {{Object}} - Results with t and y arrays
 */
function simulate(tEnd = 10.0, dt = 0.01) {{
    let y = [{init_str}];
    let t = 0;
    
    const results = {{ t: [], y: [] }};
    
    while (t < tEnd) {{
        results.t.push(t);
        results.y.push([...y]);
        y = rk4Step(t, y, dt);
        t += dt;
    }}
    
    return results;
}}

// Export for Node.js or ES6 modules
if (typeof module !== 'undefined' && module.exports) {{
    module.exports = {{ simulate, equationsOfMotion, rk4Step }};
}}

// Run if executed directly
if (typeof require !== 'undefined' && require.main === module) {{
    console.log('Running {self.system_name} simulation...');
    const results = simulate();
    console.log(`Simulation complete: ${{results.t.length}} points`);
    console.log(`Final state: ${{results.y[results.y.length-1]}}`);
}}
"""
        return template
