//! Bounded measurements, deliberately separate from correctness tests and SCORE.
use std::cell::Cell;
use std::collections::{BTreeMap, BTreeSet};
use std::io::{Read, Write};
use std::path::PathBuf;
use std::process::{Command, Stdio};
use std::rc::Rc;
use std::sync::{
    atomic::{AtomicBool, Ordering},
    Arc,
};
use std::thread;
use std::time::{Duration, Instant};

use anysystem::logger::LogEntry;
use anysystem::python::PyProcessFactory;
use anysystem::{Context, Message, Process, System};
use serde::{Deserialize, Serialize};

use crate::common::{BroadcastMessage, TestConfig};
use crate::{shared_state, Args};

const SIZES: [u64; 4] = [5, 10, 20, 50];
const SCENARIOS: [&str; 3] = ["NORMAL", "CRASHES", "DELAYED"];
const MESSAGES: usize = 10;
const HORIZON: f64 = 109.; // 100 model-time units after the last SEND at t=9.
const EVENT_LIMIT: u64 = 100_000;
const OUTPUT_LIMIT: u64 = 2 * 1024 * 1024;
const RESULT_PREFIX: &str = "OPTIMIZATION RESULT ";

// Coverage counts first valid DELIVERs to correct processes, including the
// sender. Missing deliveries at the horizon do not prove a liveness violation.
// Latencies use model time and only those deliveries. Network counters include
// service traffic; duplicate/unexpected/malformed outputs count all processes.
#[derive(Debug, Deserialize, Serialize)]
struct Measurement {
    scenario: String,
    n: u64,
    seed: u64,
    status: String,
    observed_until: f64,
    sent: usize,
    messages: u64,
    bytes: u64,
    max_sent: u64,
    delivered: usize,
    expected: usize,
    latency_mean: Option<f64>,
    latency_max: Option<f64>,
    duplicates: usize,
    unexpected: usize,
    malformed: usize,
}

// A native test driver schedules external actions independently of the student's
// traffic. It has no network traffic and is excluded from all measurements.
#[derive(Clone)]
struct Workload {
    action: Rc<Cell<usize>>,
    crashes: bool,
}

impl Process for Workload {
    fn on_local_message(&mut self, _: Message, ctx: &mut Context) -> Result<(), String> {
        for i in 1..MESSAGES {
            ctx.set_timer(&i.to_string(), i as f64);
        }
        if self.crashes {
            ctx.set_timer("10", 4.5);
        }
        ctx.set_timer("11", HORIZON);
        Ok(())
    }

    fn on_message(&mut self, _: Message, _: String, _: &mut Context) -> Result<(), String> {
        Ok(())
    }

    fn on_timer(&mut self, timer: String, _: &mut Context) -> Result<(), String> {
        self.action.set(timer.parse().unwrap());
        Ok(())
    }
}

fn send(sys: &mut System, index: usize) {
    let text = format!("{index:02}:{}", "x".repeat(61)); // fixed 64-byte payload
    sys.send_local_message(
        "0",
        Message::json("SEND", &BroadcastMessage { text: &text }),
    );
}

fn measure(config: &TestConfig, scenario: &str) -> Measurement {
    let mut sys = System::new(config.seed);
    if scenario == "DELAYED" {
        sys.network().set_delays(1., 10.);
    } else {
        sys.network().set_delays(1., 3.);
    }
    let processes: Vec<String> = (0..config.proc_count).map(|i| i.to_string()).collect();
    for proc in &processes {
        let implementation = config
            .proc_factory
            .build((proc, processes.clone()), config.seed);
        sys.add_node(proc);
        sys.add_process(proc, Box::new(implementation), proc);
    }
    let action = Rc::new(Cell::new(0));
    sys.add_node("_workload");
    sys.add_process(
        "_workload",
        Box::new(Workload {
            action: action.clone(),
            crashes: scenario == "CRASHES",
        }),
        "_workload",
    );
    sys.send_local_message("_workload", Message::new("START", "{}"));
    send(&mut sys, 0);
    let mut status = "event limit";
    for _ in 0..EVENT_LIMIT {
        if !sys.step() {
            status = "queue exhausted";
            break;
        }
        match action.replace(0) {
            0 => {}
            index @ 1..=9 => send(&mut sys, index),
            10 => {
                sys.crash_node("1");
                sys.crash_node("2");
            }
            11 => {
                status = "horizon";
                break;
            }
            _ => unreachable!(),
        }
    }
    summarize(&sys, config, scenario, status)
}

fn summarize(sys: &System, config: &TestConfig, scenario: &str, status: &str) -> Measurement {
    let mut latencies = Vec::new();
    let mut duplicates = 0;
    let mut unexpected = 0;
    let mut malformed = 0;
    let mut network_sends = BTreeMap::<String, u64>::new();
    let mut sent = BTreeMap::new();
    let mut seen = BTreeSet::new();
    // Trace order matters: a guessed DELIVER before SEND is unexpected even
    // when the two events have the same model timestamp.
    for entry in sys.logger().trace() {
        match entry {
            LogEntry::MessageSent {
                src_proc, dst_proc, ..
            } if src_proc != dst_proc => {
                *network_sends.entry(src_proc.clone()).or_default() += 1;
            }
            LogEntry::LocalMessageReceived {
                time, proc, msg, ..
            } if proc == "0" && msg.tip == "SEND" => {
                let data: serde_json::Value = serde_json::from_str(&msg.data).unwrap();
                sent.insert(data["text"].as_str().unwrap().to_string(), *time);
            }
            LogEntry::LocalMessageSent {
                time, proc, msg, ..
            } => {
                let data = serde_json::from_str::<serde_json::Value>(&msg.data).ok();
                let text = data.as_ref().and_then(|v| v["text"].as_str());
                if msg.tip != "DELIVER" || text.is_none() {
                    malformed += 1;
                    continue;
                }
                let text = text.unwrap();
                if !seen.insert((proc.clone(), text.to_string())) {
                    duplicates += 1;
                    continue;
                }
                if let Some(sent_at) = sent.get(text) {
                    if scenario != "CRASHES" || (proc != "1" && proc != "2") {
                        latencies.push(time - sent_at);
                    }
                } else {
                    unexpected += 1;
                }
            }
            _ => {}
        }
    }
    let correct_count = config.proc_count - if scenario == "CRASHES" { 2 } else { 0 };
    let network = sys.network();
    Measurement {
        scenario: scenario.to_string(),
        n: config.proc_count,
        seed: config.seed,
        status: status.to_string(),
        observed_until: sys.time(),
        sent: sent.len(),
        messages: network.network_message_count(),
        bytes: network.traffic(),
        max_sent: network_sends.into_values().max().unwrap_or(0),
        delivered: latencies.len(),
        expected: MESSAGES * correct_count as usize,
        latency_mean: (!latencies.is_empty())
            .then(|| latencies.iter().sum::<f64>() / latencies.len() as f64),
        latency_max: latencies.into_iter().reduce(f64::max),
        duplicates,
        unexpected,
        malformed,
    }
}

pub fn worker(args: &Args) {
    // This process may panic, hang in Python, or print SCORE from the validator.
    // The supervisor captures its output and never forwards it verbatim.
    if !shared_state::validate_or_report(&args.solution_path, &["BroadcastProcess"]) {
        return;
    }
    let factory = PyProcessFactory::new(&args.solution_path, "BroadcastProcess");
    let mut rows = Vec::new();
    'measurements: for scenario in SCENARIOS {
        for n in SIZES {
            for run in 0..args.optimization_runs {
                let config = TestConfig {
                    proc_factory: &factory,
                    proc_count: n,
                    seed: args.seed.wrapping_add(run as u64),
                    monkeys: 0,
                    debug: false,
                };
                let result = measure(&config, scenario);
                let incomplete = result.status == "event limit";
                rows.push(result);
                if incomplete {
                    break 'measurements;
                }
            }
        }
    }
    drop(factory);
    // Emit one final record only after all simulator callbacks and cleanup.
    // Earlier stdout belongs to the solution and is diagnostic text, even if it
    // resembles this record. This framing is not a sandbox for hostile Python.
    println!("\n{RESULT_PREFIX}{}", serde_json::to_string(&rows).unwrap());
    std::io::stdout().flush().unwrap();
}

struct WorkerResult {
    rows: Vec<Measurement>,
    problem: Option<String>,
    diagnostics: Vec<String>,
}

fn read_bounded(
    reader: impl Read + Send + 'static,
    exceeded: Arc<AtomicBool>,
) -> thread::JoinHandle<Vec<u8>> {
    thread::spawn(move || {
        let mut bytes = Vec::new();
        let _ = reader.take(OUTPUT_LIMIT + 1).read_to_end(&mut bytes);
        if bytes.len() as u64 > OUTPUT_LIMIT {
            exceeded.store(true, Ordering::Relaxed);
        }
        bytes
    })
}

fn run_worker(args: &Args, path: &str, deadline: Instant) -> std::io::Result<WorkerResult> {
    if Instant::now() >= deadline {
        return Ok(WorkerResult {
            rows: Vec::new(),
            problem: Some("shared wall-clock budget exhausted".to_string()),
            diagnostics: Vec::new(),
        });
    }
    let mut child = Command::new(std::env::current_exe()?)
        .args([
            "--optimization-worker",
            "-i",
            path,
            "-s",
            &args.seed.to_string(),
            "--optimization-runs",
            &args.optimization_runs.to_string(),
        ])
        .env("PYTHONUNBUFFERED", "1")
        // The validator also initializes Python, so set the seed before spawn.
        .env("PYTHONHASHSEED", args.seed.to_string())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;
    let exceeded = Arc::new(AtomicBool::new(false));
    let out = read_bounded(child.stdout.take().unwrap(), exceeded.clone());
    let err = read_bounded(child.stderr.take().unwrap(), exceeded.clone());
    let mut problem = None;
    loop {
        match child.try_wait() {
            Ok(Some(status)) => {
                if !status.success() {
                    problem = Some(format!("worker exited with {status}"));
                }
                break;
            }
            Ok(None) => {}
            Err(error) => {
                problem = Some(error.to_string());
                let _ = child.kill();
                let _ = child.wait();
                break;
            }
        }
        if exceeded.load(Ordering::Relaxed) {
            problem = Some("output limit".to_string());
            let _ = child.kill();
            let _ = child.wait();
            break;
        }
        if Instant::now() >= deadline {
            problem = Some(format!(
                "shared wall-clock limit ({}s)",
                args.optimization_timeout
            ));
            let _ = child.kill();
            let _ = child.wait();
            break;
        }
        thread::sleep(Duration::from_millis(10));
    }
    let stdout = out.join().unwrap_or_default();
    let stderr = err.join().unwrap_or_default();
    if stdout.len() as u64 > OUTPUT_LIMIT || stderr.len() as u64 > OUTPUT_LIMIT {
        problem = Some("output limit".to_string());
    }
    let output = String::from_utf8_lossy(&stdout);
    let (diagnostic_output, result_line) = output
        .strip_suffix('\n')
        .unwrap_or(&output)
        .rsplit_once('\n')
        .unwrap_or(("", output.as_ref()));
    let mut rows = Vec::new();
    if problem.is_none() {
        match parse_measurements(result_line.trim_end_matches('\r'), args) {
            Ok(measurements) => rows = measurements,
            Err(error) => problem = Some(error),
        }
    }
    let mut diagnostics = Vec::new();
    for line in diagnostic_output.lines() {
        if !line.is_empty() && diagnostics.len() < 4 {
            diagnostics.push(line.chars().take(240).collect());
        }
    }
    // A missing or malformed final record is also useful diagnostic output.
    if rows.is_empty() && !result_line.is_empty() && diagnostics.len() < 4 {
        diagnostics.push(result_line.chars().take(240).collect());
    }
    for line in String::from_utf8_lossy(&stderr).lines().take(4) {
        diagnostics.push(line.chars().take(240).collect());
    }
    Ok(WorkerResult {
        rows,
        problem,
        diagnostics,
    })
}

fn parse_measurements(line: &str, args: &Args) -> Result<Vec<Measurement>, String> {
    let json = line
        .strip_prefix(RESULT_PREFIX)
        .ok_or("worker did not return measurements")?;
    let rows: Vec<Measurement> =
        serde_json::from_str(json).map_err(|e| format!("invalid measurements: {e}"))?;
    let expected = SCENARIOS.iter().flat_map(|scenario| {
        SIZES.iter().flat_map(move |n| {
            (0..args.optimization_runs)
                .map(move |run| (*scenario, *n, args.seed.wrapping_add(run as u64)))
        })
    });
    let total = SCENARIOS.len() * SIZES.len() * args.optimization_runs as usize;
    if rows.is_empty()
        || rows.len() > total
        || rows.iter().zip(expected).any(|(row, (scenario, n, seed))| {
            row.scenario != scenario || row.n != n || row.seed != seed
        })
        || rows[..rows.len() - 1]
            .iter()
            .any(|row| row.status == "event limit")
        || (rows.len() < total && rows.last().unwrap().status != "event limit")
    {
        return Err("worker returned an incomplete or unexpected measurement sequence".to_string());
    }
    Ok(rows)
}

fn optimized_path(args: &Args) -> PathBuf {
    args.optimized_impl
        .as_ref()
        .map(PathBuf::from)
        .unwrap_or_else(|| {
            PathBuf::from(&args.solution_path).with_file_name("broadcast_scalable.py")
        })
}

pub fn run_if_present(args: &Args, score: f32) {
    if args.optimized_impl.is_some() || optimized_path(args).exists() {
        if score == 7. {
            run(args);
        } else {
            println!("SKIPPED: OPTIMIZATION (requires main SCORE: 7)");
        }
    }
}

fn cell(row: Option<&Measurement>) -> String {
    let Some(r) = row else {
        return "NOT MEASURED".to_string();
    };
    let latency = |value: Option<f64>| {
        value
            .map(|v| format!("{v:.3}"))
            .unwrap_or_else(|| "n/a".to_string())
    };
    format!(
        "{} {} {} {}/{} {} {} {} {} {} [{}; t={:.3}; sent={}]",
        r.messages,
        r.bytes,
        r.max_sent,
        r.delivered,
        r.expected,
        latency(r.latency_mean),
        latency(r.latency_max),
        r.duplicates,
        r.unexpected,
        r.malformed,
        if r.status == "event limit" {
            "INCOMPLETE: event limit"
        } else {
            &r.status
        },
        r.observed_until,
        r.sent
    )
}

pub fn run(args: &Args) {
    println!("\n--- OPTIMIZATION ---\n");
    let optimized = optimized_path(args);
    if !optimized.is_file() {
        println!(
            "SKIPPED: scalable implementation not found: {}",
            optimized.display()
        );
        return;
    }
    // Both implementations receive the same external workload: ten unique
    // 64-byte messages from process 0 at t=0..9, observed until HORIZON.
    // Run indices start at zero; internal traffic depends on the implementation.
    // Measurements have no performance threshold or automatic optimization score.
    // Event/time/output limits mark incomplete measurements without changing SCORE.
    println!("Workload: N={SIZES:?}, messages={MESSAGES} x 64B, sender=0, send_t=0..9, horizon={HORIZON}");
    println!("Delays: NORMAL/CRASHES=[1,3], DELAYED=[1,10]; CRASHES: [1,2] at t=4.5");
    println!("Runs={}, base_seed={}; limits: {EVENT_LIMIT} events/case, {}s/comparison, {OUTPUT_LIMIT} bytes/stream", args.optimization_runs, args.seed, args.optimization_timeout);
    let deadline = Instant::now() + Duration::from_secs(args.optimization_timeout);
    let mut results = Vec::new();
    for (label, path) in [
        ("BASE", args.solution_path.as_str()),
        ("SCALABLE", optimized.to_str().unwrap_or("")),
    ] {
        if label == "SCALABLE"
            && results.first().is_some_and(|r: &WorkerResult| {
                r.problem.is_some() || r.rows.iter().any(|row| row.status == "event limit")
            })
        {
            println!("SKIPPED: SCALABLE (BASE measurement incomplete)");
            results.push(WorkerResult {
                rows: Vec::new(),
                problem: None,
                diagnostics: Vec::new(),
            });
            continue;
        }
        println!("Measuring {label}: {path}");
        let result = run_worker(args, path, deadline).unwrap_or_else(|e| WorkerResult {
            rows: Vec::new(),
            problem: Some(e.to_string()),
            diagnostics: Vec::new(),
        });
        if let Some(problem) = &result.problem {
            println!("{label}: INCOMPLETE: {problem}");
        }
        for line in &result.diagnostics {
            println!("{label} diagnostic: {line}");
        }
        results.push(result);
    }
    println!("scenario N seed | BASE: messages bytes max_sent delivered/expected latency_mean latency_max dup unexpected malformed | SCALABLE: same columns");
    for scenario in SCENARIOS {
        for n in SIZES {
            for run in 0..args.optimization_runs {
                let seed = args.seed.wrapping_add(run as u64);
                let find = |side: usize| {
                    results[side]
                        .rows
                        .iter()
                        .find(|r| r.scenario == scenario && r.n == n && r.seed == seed)
                };
                println!(
                    "{scenario} {n} {seed} | {} | {}",
                    cell(find(0)),
                    cell(find(1))
                );
            }
        }
    }
}