use crate::ir::*;
use std::fmt::Write;

/// Generate complete Rust source from the IR.
pub fn generate(ir: &SeqIR) -> String {
    let mut out = String::with_capacity(4096);
    gen_header(&mut out, ir);
    gen_constants(&mut out, ir);
    gen_vars_struct(&mut out, ir);
    gen_program_vars_impl(&mut out, ir);
    gen_meta_impl(&mut out, ir);
    gen_state_sets(&mut out, ir);
    gen_main(&mut out, ir);
    out
}

fn gen_header(out: &mut String, _ir: &SeqIR) {
    writeln!(out, "//! Generated by snc — do not edit.").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "use epics_seq::prelude::*;").unwrap();
    writeln!(out).unwrap();
}

fn gen_constants(out: &mut String, ir: &SeqIR) {
    // Channel ID constants
    for ch in &ir.channels {
        let name = ch.var_name.to_uppercase();
        writeln!(out, "const CH_{name}: usize = {};", ch.id).unwrap();
    }
    if !ir.channels.is_empty() {
        writeln!(out).unwrap();
    }

    // Event flag ID constants
    for ef in &ir.event_flags {
        let name = ef.name.to_uppercase();
        writeln!(out, "const EF_{name}: usize = {};", ef.id).unwrap();
    }
    if !ir.event_flags.is_empty() {
        writeln!(out).unwrap();
    }

    // State ID constants per state set
    for ss in &ir.state_sets {
        let ss_upper = ss.name.to_uppercase();
        for state in &ss.states {
            let st_upper = state.name.to_uppercase();
            writeln!(out, "const {ss_upper}_{st_upper}: usize = {};", state.id).unwrap();
        }
    }
    if !ir.state_sets.is_empty() {
        writeln!(out).unwrap();
    }
}

fn gen_vars_struct(out: &mut String, ir: &SeqIR) {
    let name = &ir.program_name;
    let struct_name = format!("{name}Vars");

    writeln!(out, "#[derive(Clone)]").unwrap();
    writeln!(out, "struct {struct_name} {{").unwrap();
    for var in &ir.variables {
        let rust_type = var.var_type.rust_type();
        writeln!(out, "    {}: {rust_type},", var.name).unwrap();
    }
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn gen_program_vars_impl(out: &mut String, ir: &SeqIR) {
    let struct_name = format!("{}Vars", ir.program_name);

    writeln!(out, "impl ProgramVars for {struct_name} {{").unwrap();

    // get_channel_value
    writeln!(
        out,
        "    fn get_channel_value(&self, ch_id: usize) -> EpicsValue {{"
    )
    .unwrap();
    writeln!(out, "        match ch_id {{").unwrap();
    for var in &ir.variables {
        if let Some(ch_id) = var.channel_id {
            let expr = var.var_type.to_epics_value_expr(&format!("self.{}", var.name));
            writeln!(out, "            {ch_id} => {expr},").unwrap();
        }
    }
    writeln!(out, "            _ => EpicsValue::Double(0.0),").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();

    // set_channel_value
    writeln!(
        out,
        "    fn set_channel_value(&mut self, ch_id: usize, value: &EpicsValue) {{"
    )
    .unwrap();
    writeln!(out, "        match ch_id {{").unwrap();
    for var in &ir.variables {
        if let Some(ch_id) = var.channel_id {
            let expr = var.var_type.from_epics_value_expr("value");
            writeln!(out, "            {ch_id} => self.{} = {expr},", var.name).unwrap();
        }
    }
    writeln!(out, "            _ => {{}}").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn gen_meta_impl(out: &mut String, ir: &SeqIR) {
    let struct_name = format!("{}Meta", ir.program_name);
    let vars_name = format!("{}Vars", ir.program_name);

    writeln!(out, "struct {struct_name};").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "impl ProgramMeta for {struct_name} {{").unwrap();
    writeln!(
        out,
        "    const NUM_CHANNELS: usize = {};",
        ir.channels.len()
    )
    .unwrap();
    writeln!(
        out,
        "    const NUM_EVENT_FLAGS: usize = {};",
        ir.event_flags.len()
    )
    .unwrap();
    writeln!(
        out,
        "    const NUM_STATE_SETS: usize = {};",
        ir.state_sets.len()
    )
    .unwrap();
    writeln!(out).unwrap();

    // channel_defs()
    writeln!(out, "    fn channel_defs() -> Vec<ChannelDef> {{").unwrap();
    writeln!(out, "        vec![").unwrap();
    for ch in &ir.channels {
        let monitored = ch.monitored;
        let sync_ef = match ch.sync_ef {
            Some(id) => format!("Some({id})"),
            None => "None".to_string(),
        };
        writeln!(out, "            ChannelDef {{").unwrap();
        writeln!(out, "                var_name: \"{}\".into(),", ch.var_name).unwrap();
        writeln!(out, "                pv_name: \"{}\".into(),", ch.pv_name).unwrap();
        writeln!(out, "                monitored: {monitored},").unwrap();
        writeln!(out, "                sync_ef: {sync_ef},").unwrap();
        writeln!(out, "            }},").unwrap();
    }
    writeln!(out, "        ]").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();

    // event_flag_sync_map()
    writeln!(out, "    fn event_flag_sync_map() -> Vec<Vec<usize>> {{").unwrap();
    writeln!(out, "        vec![").unwrap();
    for ef in &ir.event_flags {
        let chs: Vec<String> = ef.synced_channels.iter().map(|c| c.to_string()).collect();
        writeln!(out, "            vec![{}],", chs.join(", ")).unwrap();
    }
    writeln!(out, "        ]").unwrap();
    writeln!(out, "    }}").unwrap();

    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    // Make the vars struct name and meta struct name available
    // (they're referenced by gen_main)
    let _ = vars_name;
}

fn gen_state_sets(out: &mut String, ir: &SeqIR) {
    for ss in &ir.state_sets {
        gen_one_state_set(out, ir, ss);
    }
}

fn gen_one_state_set(out: &mut String, ir: &SeqIR, ss: &IRStateSet) {
    let struct_name = format!("{}Vars", ir.program_name);
    let fn_name = &ss.name;
    let state_names: Vec<&str> = ss.states.iter().map(|s| s.name.as_str()).collect();

    writeln!(
        out,
        "async fn {fn_name}(mut ctx: StateSetContext<{struct_name}>) -> SeqResult<()> {{"
    )
    .unwrap();

    // State name array for debugging
    let names_str: Vec<String> = state_names.iter().map(|n| format!("\"{n}\"")).collect();
    writeln!(
        out,
        "    let _state_names = [{}];",
        names_str.join(", ")
    )
    .unwrap();
    writeln!(out).unwrap();

    // SS-local variables
    for var in &ss.local_vars {
        let default_val = var.var_type.default_value();
        let init = var
            .init_value
            .as_deref()
            .unwrap_or(&default_val);
        writeln!(
            out,
            "    let mut {}: {} = {init};",
            var.name,
            var.var_type.rust_type()
        )
        .unwrap();
    }
    if !ss.local_vars.is_empty() {
        writeln!(out).unwrap();
    }

    writeln!(out, "    ctx.enter_state(0);").unwrap();
    writeln!(out, "    ctx.wakeup().notify_one();").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "    loop {{").unwrap();
    writeln!(out, "        if ctx.is_shutdown() {{ break; }}").unwrap();
    writeln!(out).unwrap();

    // Entry block per state
    writeln!(out, "        if ctx.should_run_entry() {{").unwrap();
    writeln!(out, "            match ctx.current_state() {{").unwrap();
    for state in &ss.states {
        if let Some(entry) = &state.entry {
            let ss_upper = ss.name.to_uppercase();
            let st_upper = state.name.to_uppercase();
            writeln!(out, "                {ss_upper}_{st_upper} => {{").unwrap();
            for line in entry.code.lines() {
                writeln!(out, "                    {line}").unwrap();
            }
            writeln!(out, "                }}").unwrap();
        }
    }
    writeln!(out, "                _ => {{}}").unwrap();
    writeln!(out, "            }}").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out).unwrap();

    // Inner loop
    writeln!(out, "        loop {{").unwrap();
    writeln!(out, "            ctx.wait_for_wakeup().await;").unwrap();
    writeln!(out, "            if ctx.is_shutdown() {{ return Ok(()); }}").unwrap();
    writeln!(out, "            ctx.reset_wakeup();").unwrap();
    writeln!(out, "            ctx.sync_dirty_vars();").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "            match ctx.current_state() {{").unwrap();
    for state in &ss.states {
        let ss_upper = ss.name.to_uppercase();
        let st_upper = state.name.to_uppercase();
        writeln!(out, "                {ss_upper}_{st_upper} => {{").unwrap();

        for (i, trans) in state.transitions.iter().enumerate() {
            let kw = if i == 0 { "if" } else { "} else if" };
            let cond = trans.condition.as_deref().unwrap_or("true");
            writeln!(out, "                    {kw} {cond} {{").unwrap();

            // Action code
            for line in trans.action.code.lines() {
                writeln!(out, "                        {line}").unwrap();
            }

            // Transition
            match trans.target_state {
                Some(target) => {
                    // Find target state name
                    let target_name = ss.states.get(target).map_or("?", |s| &s.name);
                    let target_upper = target_name.to_uppercase();
                    writeln!(
                        out,
                        "                        ctx.transition_to({ss_upper}_{target_upper});"
                    )
                    .unwrap();
                }
                None => {
                    writeln!(out, "                        return Ok(());").unwrap();
                }
            }
        }
        if !state.transitions.is_empty() {
            writeln!(out, "                    }}").unwrap();
        }
        writeln!(out, "                }}").unwrap();
    }
    writeln!(out, "                _ => return Err(SeqError::InvalidStateId(ctx.current_state())),").unwrap();
    writeln!(out, "            }}").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "            if ctx.has_transition() {{ break; }}").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out).unwrap();

    // Exit block per state
    writeln!(out, "        if ctx.should_run_exit() {{").unwrap();
    writeln!(out, "            match ctx.current_state() {{").unwrap();
    for state in &ss.states {
        if let Some(exit) = &state.exit {
            let ss_upper = ss.name.to_uppercase();
            let st_upper = state.name.to_uppercase();
            writeln!(out, "                {ss_upper}_{st_upper} => {{").unwrap();
            for line in exit.code.lines() {
                writeln!(out, "                    {line}").unwrap();
            }
            writeln!(out, "                }}").unwrap();
        }
    }
    writeln!(out, "                _ => {{}}").unwrap();
    writeln!(out, "            }}").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "        if let Some(next) = ctx.take_transition() {{").unwrap();
    writeln!(out, "            ctx.enter_state(next);").unwrap();
    writeln!(out, "            ctx.wakeup().notify_one();").unwrap();
    writeln!(out, "        }}").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "    Ok(())").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn gen_main(out: &mut String, ir: &SeqIR) {
    let name = &ir.program_name;
    let vars_name = format!("{name}Vars");
    let meta_name = format!("{name}Meta");

    writeln!(out, "#[tokio::main]").unwrap();
    writeln!(
        out,
        "async fn main() -> Result<(), Box<dyn std::error::Error>> {{"
    )
    .unwrap();
    writeln!(out, "    tracing_subscriber::fmt::init();").unwrap();
    writeln!(out).unwrap();
    writeln!(
        out,
        "    let macro_str = std::env::args().nth(1).unwrap_or_default();"
    )
    .unwrap();
    writeln!(out).unwrap();

    // Initial variable values
    writeln!(out, "    let initial = {vars_name} {{").unwrap();
    for var in &ir.variables {
        let default_val = var.var_type.default_value();
        let init = var
            .init_value
            .as_deref()
            .unwrap_or(&default_val);
        writeln!(out, "        {}: {init},", var.name).unwrap();
    }
    writeln!(out, "    }};").unwrap();
    writeln!(out).unwrap();

    writeln!(
        out,
        "    ProgramBuilder::<{vars_name}, {meta_name}>::new(\"{name}\", initial)"
    )
    .unwrap();
    writeln!(out, "        .macros(&macro_str)").unwrap();
    for ss in &ir.state_sets {
        let fn_name = &ss.name;
        writeln!(
            out,
            "        .add_ss(Box::new(|ctx| Box::pin({fn_name}(ctx))))"
        )
        .unwrap();
    }
    writeln!(out, "        .run()").unwrap();
    writeln!(out, "        .await?;").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "    Ok(())").unwrap();
    writeln!(out, "}}").unwrap();
}
