"""
Rust Code Generator for MechanicsDSL

Generates standalone Rust simulation code using ode_solvers crate.
"""

from sympy.printing.rust import RustCodePrinter

from ..utils import logger
from .base import CodeGenerator


def rust_code(expr):
    """Convert SymPy expression to Rust code string."""
    printer = RustCodePrinter()
    return printer.doprint(expr)


class RustGenerator(CodeGenerator):
    """
    Generates Rust simulation code with ode_solvers crate.

    Produces standalone binaries with CSV output.
    """

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

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

    def generate(self, output_file: str = "simulation.rs") -> str:
        """Generate Rust simulation code."""
        logger.info(f"Generating Rust 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_cargo_toml(self, output_dir: str = ".", embedded: bool = False) -> str:
        """Generate Cargo.toml for the Rust project.

        Args:
            output_dir: Directory to write Cargo.toml
            embedded: If True, generate no_std compatible config
        """
        import os

        cargo_path = os.path.join(output_dir, "Cargo.toml")

        if embedded:
            # Embedded/no_std configuration
            cargo_content = f"""[package]
name = "{self.system_name.replace("-", "_")}"
version = "0.1.0"
edition = "2021"
authors = ["MechanicsDSL Generator"]
description = "Physics simulation generated by MechanicsDSL (embedded)"

[dependencies]
# No standard library dependencies for embedded targets
libm = "0.2"  # Math functions for no_std

[profile.release]
opt-level = 3
lto = true
codegen-units = 1
panic = "abort"

# ARM Cortex-M targets
# Build with: cargo build --target thumbv7em-none-eabihf --release
"""
        else:
            # Standard configuration
            cargo_content = f"""[package]
name = "{self.system_name.replace("-", "_")}"
version = "0.1.0"
edition = "2021"
authors = ["MechanicsDSL Generator"]
description = "Physics simulation generated by MechanicsDSL"

[dependencies]
# High-performance ODE solvers
# Note: ode_solvers is optional if using built-in RK4

[profile.release]
opt-level = 3
lto = true

[profile.release-arm]
inherits = "release"
# ARM-specific optimizations
# Build with: cargo build --profile release-arm --target aarch64-unknown-linux-gnu
"""

        with open(cargo_path, "w") as f:
            f.write(cargo_content)

        logger.info(f"Generated Cargo.toml at {cargo_path}")
        return cargo_path

    def generate_project(self, output_dir: str = ".", embedded: bool = False) -> dict:
        """Generate complete Rust project with Cargo structure.

        Args:
            output_dir: Root directory for the project
            embedded: If True, generate no_std compatible project

        Returns:
            dict of generated file paths
        """
        import os

        # Create project structure
        src_dir = os.path.join(output_dir, "src")
        os.makedirs(src_dir, exist_ok=True)

        # Generate main.rs
        main_rs = os.path.join(src_dir, "main.rs")
        self.generate(main_rs)

        # Generate Cargo.toml
        cargo_toml = self.generate_cargo_toml(output_dir, embedded)

        # Generate README
        readme_path = os.path.join(output_dir, "README.md")
        target_note = "ARM Cortex-M embedded target" if embedded else "standard targets"
        readme_content = f"""# {self.system_name}

Auto-generated Rust physics simulation by MechanicsDSL.

## Build & Run

### Standard Build
```bash
cargo run --release
```

### Cross-compile for Raspberry Pi (aarch64)
```bash
rustup target add aarch64-unknown-linux-gnu
cargo build --release --target aarch64-unknown-linux-gnu
```

### Cross-compile for ARM Cortex-M (embedded)
```bash
rustup target add thumbv7em-none-eabihf
cargo build --release --target thumbv7em-none-eabihf
```

## Output

Results are saved to `{self.system_name}.csv`.

## Generated for: {target_note}
"""
        with open(readme_path, "w") as f:
            f.write(readme_content)

        logger.info(f"Generated complete Rust project in {output_dir}")
        return {"main": main_rs, "cargo": cargo_toml, "readme": readme_path}

    def generate_equations(self) -> str:
        """Generate Rust 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]
                rs_expr = rust_code(expr)
                lines.append(f"        dydt[{idx+1}] = {rs_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 Rust simulation code."""
        # Parameter definitions
        param_lines = []
        for name, val in self.parameters.items():
            param_lines.append(f"const {name.upper()}: f64 = {val};")
        param_str = "\n".join(param_lines)

        # State variable unpacking
        unpack_lines = []
        idx = 0
        for coord in self.coordinates:
            unpack_lines.append(f"        let {coord} = y[{idx}];")
            unpack_lines.append(f"        let {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

        # Pre-compute CSV format strings to avoid f-string escaping issues
        csv_header = ",".join(self.coordinates)
        csv_format = ",".join(f"y[{i*2}]" for i in range(len(self.coordinates)))
        csv_placeholders = ",".join(["{}" for _ in self.coordinates])

        template = f"""//! Auto-generated simulation: {self.system_name}
//! Generated by MechanicsDSL
//! 
//! Build with: cargo build --release
//! Requires ode_solvers in Cargo.toml:
//!   [dependencies]
//!   ode_solvers = "0.4"

use std::fs::File;
use std::io::Write;

// Physical Parameters
{param_str}

const DIM: usize = {state_dim};

fn equations_of_motion(_t: f64, y: &[f64; DIM], dydt: &mut [f64; DIM]) {{
    // Unpack state
{unpack_str}
    
    // Compute derivatives
{eq_str}
}}

fn rk4_step(t: f64, y: &mut [f64; DIM], dt: f64) {{
    let mut k1 = [0.0f64; DIM];
    let mut k2 = [0.0f64; DIM];
    let mut k3 = [0.0f64; DIM];
    let mut k4 = [0.0f64; DIM];
    let mut tmp = [0.0f64; DIM];
    
    equations_of_motion(t, y, &mut k1);
    
    for i in 0..DIM {{
        tmp[i] = y[i] + 0.5 * dt * k1[i];
    }}
    equations_of_motion(t + 0.5 * dt, &tmp, &mut k2);
    
    for i in 0..DIM {{
        tmp[i] = y[i] + 0.5 * dt * k2[i];
    }}
    equations_of_motion(t + 0.5 * dt, &tmp, &mut k3);
    
    for i in 0..DIM {{
        tmp[i] = y[i] + dt * k3[i];
    }}
    equations_of_motion(t + dt, &tmp, &mut k4);
    
    for i in 0..DIM {{
        y[i] += dt / 6.0 * (k1[i] + 2.0 * k2[i] + 2.0 * k3[i] + k4[i]);
    }}
}}

fn main() {{
    println!("Running {self.system_name} simulation...");
    
    let mut y: [f64; DIM] = [{init_str}];
    let dt = 0.01;
    let t_end = 10.0;
    let mut t = 0.0;
    
    let mut file = File::create("{self.system_name}.csv").expect("Cannot create file");
    writeln!(file, "t,{csv_header}").unwrap();
    
    while t < t_end {{
        writeln!(file, "{{:.6}},{csv_placeholders}", t, {csv_format}).unwrap();
        rk4_step(t, &mut y, dt);
        t += dt;
    }}
    
    println!("Simulation complete. Results saved to {self.system_name}.csv");
}}
"""
        return template
