#![allow(missing_docs)]

use std::collections::BTreeMap;
use std::env;
use std::hint::black_box;

use telltale_machine::model::effects::{EffectFailure, EffectHandler, EffectResult};
use telltale_machine::runtime::loader::CodeImage;
use telltale_machine::{
    CommunicationReplayMode, Instr, ObservabilityRetentionConfig, ObservabilityRetentionMode,
    PayloadValidationMode, ProtocolMachine, ProtocolMachineConfig,
};
use telltale_types::{GlobalType, Label, LocalTypeR, ValType};

struct ProfileHandler;

impl EffectHandler for ProfileHandler {
    fn handle_send(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &[telltale_machine::coroutine::Value],
    ) -> EffectResult<telltale_machine::coroutine::Value> {
        EffectResult::success(telltale_machine::coroutine::Value::Unit)
    }

    fn handle_recv(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &mut Vec<telltale_machine::coroutine::Value>,
        _payload: &telltale_machine::coroutine::Value,
    ) -> EffectResult<()> {
        EffectResult::success(())
    }

    fn handle_choose(
        &self,
        _role: &str,
        _partner: &str,
        labels: &[String],
        _state: &[telltale_machine::coroutine::Value],
    ) -> EffectResult<String> {
        match labels.first().cloned() {
            Some(label) => EffectResult::success(label),
            None => EffectResult::failure(EffectFailure::invalid_input("no labels to choose from")),
        }
    }

    fn step(
        &self,
        _role: &str,
        _state: &mut Vec<telltale_machine::coroutine::Value>,
    ) -> EffectResult<()> {
        EffectResult::success(())
    }
}

fn bool_to_usize(value: bool) -> usize {
    usize::from(u8::from(value))
}

fn main() {
    let mut args = env::args().skip(1);
    let workload = args.next().unwrap_or_else(|| usage_and_exit());
    let iterations = args
        .next()
        .map(|raw| raw.parse::<usize>().unwrap_or_else(|_| usage_and_exit()))
        .unwrap_or(256);

    let checksum = match workload.as_str() {
        "scheduler-many-paused" => profile_scheduler_many_paused(iterations),
        "scheduler-many-paused-run-only" => profile_scheduler_many_paused_run_only(iterations),
        "repeated-load-reuse" => profile_repeated_load_reuse(iterations),
        "send-recv-replay-nullifier" => profile_send_recv_replay_nullifier(iterations),
        "repeated-open-same-image" => profile_repeated_open_same_image(iterations),
        _ => usage_and_exit(),
    };

    println!("profile checksum: {checksum}");
}

fn usage_and_exit() -> ! {
    eprintln!(
        "usage: cargo run -p telltale-machine --release --bin profile_driver -- \\\n\
         <scheduler-many-paused|scheduler-many-paused-run-only|repeated-load-reuse|send-recv-replay-nullifier|repeated-open-same-image> [iterations]"
    );
    std::process::exit(2);
}

fn capped_retention_config() -> ObservabilityRetentionConfig {
    ObservabilityRetentionConfig {
        mode: ObservabilityRetentionMode::Capped,
        capacity: 256,
    }
}

fn profile_scheduler_many_paused(iterations: usize) -> usize {
    let handler = ProfileHandler;
    let mut checksum = 0usize;
    for _ in 0..iterations {
        let image = yield_image(256, 8);
        let mut machine = ProtocolMachine::new(ProtocolMachineConfig {
            observability_retention: capped_retention_config(),
            ..ProtocolMachineConfig::strict_large_fanout()
        });
        machine
            .load_choreography(&image)
            .expect("load choreography");
        for idx in 1..256 {
            machine.pause_role(&format!("R{idx}"));
        }
        let status = machine.run(&handler, 10_000).expect("run machine");
        checksum ^= black_box(machine.trace().len())
            ^ black_box(bool_to_usize(matches!(
                status,
                telltale_machine::RunStatus::Stuck
            )));
    }
    checksum
}

fn profile_repeated_load_reuse(iterations: usize) -> usize {
    let image = yield_image(16, 16);
    let mut config = ProtocolMachineConfig {
        observability_retention: capped_retention_config(),
        ..ProtocolMachineConfig::strict_large_fanout()
    };
    config.max_sessions = config.max_sessions.max(iterations.saturating_add(16));
    config.max_coroutines = config
        .max_coroutines
        .max(iterations.saturating_mul(16).saturating_add(16));
    let mut machine = ProtocolMachine::new(config);
    for _ in 0..iterations {
        machine
            .load_choreography(&image)
            .expect("load choreography");
    }
    black_box(machine.memory_usage().retained_bytes.total)
}

fn profile_scheduler_many_paused_run_only(yields_per_role: usize) -> usize {
    let handler = ProfileHandler;
    let max_rounds = yields_per_role.max(1);
    let image = looping_yield_image(256);
    let mut machine = ProtocolMachine::new(ProtocolMachineConfig {
        observability_retention: capped_retention_config(),
        ..ProtocolMachineConfig::strict_large_fanout()
    });
    machine
        .load_choreography(&image)
        .expect("load choreography");
    for idx in 1..256 {
        machine.pause_role(&format!("R{idx}"));
    }
    let status = machine.run(&handler, max_rounds).expect("run machine");
    black_box(machine.trace().len())
        ^ black_box(bool_to_usize(matches!(
            status,
            telltale_machine::RunStatus::MaxRoundsExceeded
        )))
}

fn profile_send_recv_replay_nullifier(iterations: usize) -> usize {
    let handler = ProfileHandler;
    let image = typed_send_recv_image("A", "B", "msg", Some(ValType::Unit));
    let mut machine = ProtocolMachine::new(ProtocolMachineConfig {
        observability_retention: capped_retention_config(),
        communication_replay_mode: CommunicationReplayMode::Nullifier,
        payload_validation_mode: PayloadValidationMode::Structural,
        ..ProtocolMachineConfig::strict_verified()
    });
    for _ in 0..iterations {
        machine
            .load_choreography(&image)
            .expect("load choreography");
    }
    let status = machine.run(&handler, 100_000).expect("run machine");
    black_box(machine.trace().len())
        ^ black_box(bool_to_usize(matches!(
            status,
            telltale_machine::RunStatus::AllDone
        )))
}

fn profile_repeated_open_same_image(iterations: usize) -> usize {
    let image = yield_image(16, 16);
    let mut config = ProtocolMachineConfig {
        observability_retention: capped_retention_config(),
        ..ProtocolMachineConfig::strict_large_fanout()
    };
    config.max_sessions = config.max_sessions.max(iterations.saturating_add(16));
    config.max_coroutines = config
        .max_coroutines
        .max(iterations.saturating_mul(16).saturating_add(16));
    let mut machine = ProtocolMachine::new(config);
    let mut checksum = 0usize;
    for _ in 0..iterations {
        let sid = machine
            .load_choreography(&image)
            .expect("load choreography");
        checksum ^= black_box(sid) ^ black_box(machine.session_coroutines(sid).len());
    }
    checksum
}

fn yield_image(num_roles: usize, yields_per_role: usize) -> CodeImage {
    let mut programs = BTreeMap::new();
    let mut local_types = BTreeMap::new();

    for idx in 0..num_roles {
        let role = format!("R{idx}");
        let mut code = Vec::with_capacity(yields_per_role + 1);
        for _ in 0..yields_per_role {
            code.push(Instr::Yield);
        }
        code.push(Instr::Halt);

        programs.insert(role.clone(), code);
        local_types.insert(role, LocalTypeR::End);
    }

    CodeImage {
        programs,
        global_type: GlobalType::End,
        local_types,
    }
}

fn looping_yield_image(num_roles: usize) -> CodeImage {
    let mut programs = BTreeMap::new();
    let mut local_types = BTreeMap::new();

    for idx in 0..num_roles {
        let role = format!("R{idx}");
        programs.insert(role.clone(), vec![Instr::Yield, Instr::Jump { target: 0 }]);
        local_types.insert(role, LocalTypeR::End);
    }

    CodeImage {
        programs,
        global_type: GlobalType::End,
        local_types,
    }
}

fn typed_send_recv_image(
    sender: &str,
    receiver: &str,
    label: &str,
    expected: Option<ValType>,
) -> CodeImage {
    let mut local_types = BTreeMap::new();
    local_types.insert(
        sender.to_string(),
        LocalTypeR::Send {
            partner: receiver.into(),
            branches: vec![(Label::new(label), expected.clone(), LocalTypeR::End)],
        },
    );
    local_types.insert(
        receiver.to_string(),
        LocalTypeR::Recv {
            partner: sender.into(),
            branches: vec![(Label::new(label), expected, LocalTypeR::End)],
        },
    );

    let global = GlobalType::send(sender, receiver, Label::new(label), GlobalType::End);
    CodeImage::from_local_types(&local_types, &global)
}
