"""
JavaScript Code Generator for MechanicsDSL

Generates standalone JavaScript simulation code for Node.js or browsers.
"""
from typing import Dict, List
import sympy as sp
from sympy.printing.jscode import jscode

from .base import CodeGenerator
from ..utils import logger


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
