//! Shared test infrastructure for ProtocolMachine conformance tests.

use std::collections::BTreeMap;
use std::sync::Mutex;

use proptest::prelude::*;
use telltale_machine::buffer::{BackpressurePolicy, BufferConfig, BufferMode};
use telltale_machine::coroutine::Value;
use telltale_machine::model::effects::{EffectFailure, EffectHandler, EffectResult};
use telltale_machine::runtime::loader::CodeImage;
use telltale_machine::{ObsEvent, ProtocolMachine, ProtocolMachineError, StepResult};
use telltale_types::{GlobalType, Label, LocalTypeR};

/// Deterministic seed for reproducibility.
pub const SEED: [u8; 32] = [
    0x56, 0x4D, 0x43, 0x6F, 0x6E, 0x66, 0x6F, 0x72, // "PMConfor"
    0x6D, 0x61, 0x6E, 0x63, 0x65, 0x54, 0x65, 0x73, // "manceTes"
    0x74, 0x53, 0x75, 0x69, 0x74, 0x65, 0x56, 0x31, // "tSuiteV1"
    0x52, 0x75, 0x73, 0x74, 0x56, 0x4D, 0x30, 0x31, // "RustVM01"
];

// ============================================================================
// Handlers
// ============================================================================

/// Returns `Value::Nat(42)` on send, no-op on recv/step.
pub struct PassthroughHandler;

impl EffectHandler for PassthroughHandler {
    fn handle_send(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &[Value],
    ) -> EffectResult<Value> {
        EffectResult::success(Value::Nat(42))
    }

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

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

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

/// Records all (role, partner, label) triples for send/recv/step calls.
pub struct RecordingHandler {
    pub sends: Mutex<Vec<(String, String, String)>>,
    pub recvs: Mutex<Vec<(String, String, String)>>,
    pub steps: Mutex<Vec<String>>,
}

impl RecordingHandler {
    pub fn new() -> Self {
        Self {
            sends: Mutex::new(Vec::new()),
            recvs: Mutex::new(Vec::new()),
            steps: Mutex::new(Vec::new()),
        }
    }
}

impl EffectHandler for RecordingHandler {
    fn handle_send(
        &self,
        role: &str,
        partner: &str,
        label: &str,
        _state: &[Value],
    ) -> EffectResult<Value> {
        self.sends
            .lock()
            .expect("recording handler lock poisoned")
            .push((role.into(), partner.into(), label.into()));
        EffectResult::success(Value::Nat(42))
    }

    fn handle_recv(
        &self,
        role: &str,
        partner: &str,
        label: &str,
        _state: &mut Vec<Value>,
        _payload: &Value,
    ) -> EffectResult<()> {
        self.recvs
            .lock()
            .expect("recording handler lock poisoned")
            .push((role.into(), partner.into(), label.into()));
        EffectResult::success(())
    }

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

    fn step(&self, role: &str, _state: &mut Vec<Value>) -> EffectResult<()> {
        self.steps
            .lock()
            .expect("recording handler lock poisoned")
            .push(role.into());
        EffectResult::success(())
    }
}

/// Returns `Err(...)` from send/recv/step.
pub struct FailingHandler {
    pub message: String,
}

impl FailingHandler {
    pub fn new(msg: &str) -> Self {
        Self {
            message: msg.into(),
        }
    }
}

impl EffectHandler for FailingHandler {
    fn handle_send(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &[Value],
    ) -> EffectResult<Value> {
        EffectResult::failure(EffectFailure::contract_violation(self.message.clone()))
    }

    fn handle_recv(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &mut Vec<Value>,
        _payload: &Value,
    ) -> EffectResult<()> {
        EffectResult::failure(EffectFailure::contract_violation(self.message.clone()))
    }

    fn handle_choose(
        &self,
        _role: &str,
        _partner: &str,
        _labels: &[String],
        _state: &[Value],
    ) -> EffectResult<String> {
        EffectResult::failure(EffectFailure::contract_violation(self.message.clone()))
    }

    fn step(&self, _role: &str, _state: &mut Vec<Value>) -> EffectResult<()> {
        EffectResult::failure(EffectFailure::contract_violation(self.message.clone()))
    }
}

// ============================================================================
// Builders
// ============================================================================

#[derive(Debug, Clone)]
pub enum ScenarioKind {
    SimpleSendRecv,
    RecursiveSendRecv,
    Choice,
}

#[derive(Debug, Clone)]
pub struct ScenarioSpec {
    pub kind: ScenarioKind,
    pub sender: String,
    pub receiver: String,
    pub labels: Vec<String>,
}

impl ScenarioSpec {
    pub fn simple(sender: &str, receiver: &str, label: &str) -> Self {
        Self {
            kind: ScenarioKind::SimpleSendRecv,
            sender: sender.to_string(),
            receiver: receiver.to_string(),
            labels: vec![label.to_string()],
        }
    }

    pub fn recursive(sender: &str, receiver: &str, label: &str) -> Self {
        Self {
            kind: ScenarioKind::RecursiveSendRecv,
            sender: sender.to_string(),
            receiver: receiver.to_string(),
            labels: vec![label.to_string()],
        }
    }

    pub fn choice(sender: &str, receiver: &str, labels: &[&str]) -> Self {
        Self {
            kind: ScenarioKind::Choice,
            sender: sender.to_string(),
            receiver: receiver.to_string(),
            labels: labels.iter().map(|label| (*label).to_string()).collect(),
        }
    }

    pub fn to_code_image(&self) -> CodeImage {
        match self.kind {
            ScenarioKind::SimpleSendRecv => {
                simple_send_recv_image(&self.sender, &self.receiver, &self.labels[0])
            }
            ScenarioKind::RecursiveSendRecv => {
                recursive_send_recv_image(&self.sender, &self.receiver, &self.labels[0])
            }
            ScenarioKind::Choice => {
                let labels: Vec<_> = self.labels.iter().map(String::as_str).collect();
                choice_image(&self.sender, &self.receiver, &labels)
            }
        }
    }
}

/// Simple A→B:label, then End.
pub fn simple_send_recv_image(sender: &str, receiver: &str, label: &str) -> CodeImage {
    let mut local_types = BTreeMap::new();
    local_types.insert(
        sender.to_string(),
        LocalTypeR::Send {
            partner: receiver.into(),
            branches: vec![(Label::new(label), None, LocalTypeR::End)],
        },
    );
    local_types.insert(
        receiver.to_string(),
        LocalTypeR::Recv {
            partner: sender.into(),
            branches: vec![(Label::new(label), None, LocalTypeR::End)],
        },
    );

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

/// Recursive mu loop: A→B:label, B→A:label, repeat.
pub fn recursive_send_recv_image(sender: &str, receiver: &str, label: &str) -> CodeImage {
    let mut local_types = BTreeMap::new();
    local_types.insert(
        sender.to_string(),
        LocalTypeR::mu(
            "loop",
            LocalTypeR::Send {
                partner: receiver.into(),
                branches: vec![(
                    Label::new(label),
                    None,
                    LocalTypeR::Recv {
                        partner: receiver.into(),
                        branches: vec![(Label::new(label), None, LocalTypeR::var("loop"))],
                    },
                )],
            },
        ),
    );
    local_types.insert(
        receiver.to_string(),
        LocalTypeR::mu(
            "loop",
            LocalTypeR::Recv {
                partner: sender.into(),
                branches: vec![(
                    Label::new(label),
                    None,
                    LocalTypeR::Send {
                        partner: sender.into(),
                        branches: vec![(Label::new(label), None, LocalTypeR::var("loop"))],
                    },
                )],
            },
        ),
    );

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

/// Multi-branch choice: sender chooses among labels, receiver offers.
pub fn choice_image(sender: &str, receiver: &str, labels: &[&str]) -> CodeImage {
    let send_branches: Vec<_> = labels
        .iter()
        .map(|l| (Label::new(*l), None, LocalTypeR::End))
        .collect();
    let recv_branches: Vec<_> = labels
        .iter()
        .map(|l| (Label::new(*l), None, LocalTypeR::End))
        .collect();

    let mut local_types = BTreeMap::new();
    local_types.insert(
        sender.to_string(),
        LocalTypeR::send_choice(receiver, send_branches),
    );
    local_types.insert(
        receiver.to_string(),
        LocalTypeR::recv_choice(sender, recv_branches),
    );

    let global_branches: Vec<_> = labels
        .iter()
        .map(|l| (Label::new(*l), GlobalType::End))
        .collect();
    let global = GlobalType::comm(sender, receiver, global_branches);

    CodeImage::from_local_types(&local_types, &global)
}

/// Step a ProtocolMachine to completion, collecting the trace.
pub fn run_to_completion(
    machine: &mut ProtocolMachine,
    handler: &dyn EffectHandler,
    max_steps: usize,
) -> Result<Vec<ObsEvent>, ProtocolMachineError> {
    for _ in 0..max_steps {
        match machine.step(handler)? {
            StepResult::AllDone | StepResult::Stuck => break,
            StepResult::Continue => {}
        }
    }
    Ok(machine.trace().to_vec())
}

// ============================================================================
// Proptest Strategies
// ============================================================================

pub fn label_strategy() -> impl Strategy<Value = Label> {
    prop_oneof![
        Just(Label::new("msg")),
        Just(Label::new("ack")),
        Just(Label::new("data")),
        Just(Label::new("req")),
        Just(Label::new("resp")),
        Just(Label::new("yes")),
        Just(Label::new("no")),
        Just(Label::new("done")),
    ]
}

pub fn role_pair_strategy() -> impl Strategy<Value = (String, String)> {
    let roles = ["A", "B", "C", "D"];
    (0..roles.len(), 0..roles.len())
        .prop_filter("distinct roles", |(a, b)| a != b)
        .prop_map(move |(a, b)| (roles[a].to_string(), roles[b].to_string()))
}

pub fn value_strategy() -> impl Strategy<Value = Value> {
    prop_oneof![
        Just(Value::Unit),
        any::<u64>().prop_map(Value::Nat),
        any::<bool>().prop_map(Value::Bool),
        Just(Value::Str("msg".into())),
    ]
}

pub fn well_formed_global_strategy(depth: usize) -> BoxedStrategy<GlobalType> {
    if depth == 0 {
        Just(GlobalType::End).boxed()
    } else {
        prop_oneof![
            Just(GlobalType::End),
            // Simple send
            role_pair_strategy().prop_flat_map(move |(s, r)| {
                label_strategy().prop_flat_map(move |l| {
                    let s = s.clone();
                    let r = r.clone();
                    well_formed_global_strategy(depth - 1)
                        .prop_map(move |cont| GlobalType::send(&s, &r, l.clone(), cont))
                })
            }),
            // Binary choice
            role_pair_strategy().prop_flat_map(move |(s, r)| {
                (
                    well_formed_global_strategy(depth - 1),
                    well_formed_global_strategy(depth - 1),
                )
                    .prop_map(move |(c1, c2)| {
                        GlobalType::comm(
                            &s,
                            &r,
                            vec![(Label::new("yes"), c1), (Label::new("no"), c2)],
                        )
                    })
            }),
            // Guarded recursion
            role_pair_strategy().prop_map(|(s, r)| {
                GlobalType::mu(
                    "t",
                    GlobalType::comm(
                        &s,
                        &r,
                        vec![
                            (Label::new("continue"), GlobalType::var("t")),
                            (Label::new("stop"), GlobalType::End),
                        ],
                    ),
                )
            }),
        ]
        .boxed()
    }
}

pub fn buffer_config_strategy() -> impl Strategy<Value = BufferConfig> {
    prop_oneof![
        Just(BufferConfig {
            mode: BufferMode::Fifo,
            initial_capacity: 4,
            policy: BackpressurePolicy::Block,
        }),
        Just(BufferConfig {
            mode: BufferMode::Fifo,
            initial_capacity: 4,
            policy: BackpressurePolicy::Drop,
        }),
        Just(BufferConfig {
            mode: BufferMode::Fifo,
            initial_capacity: 4,
            policy: BackpressurePolicy::Error,
        }),
        Just(BufferConfig {
            mode: BufferMode::Fifo,
            initial_capacity: 4,
            policy: BackpressurePolicy::Resize { max_capacity: 64 },
        }),
        Just(BufferConfig {
            mode: BufferMode::LatestValue,
            initial_capacity: 1,
            policy: BackpressurePolicy::Block,
        }),
    ]
}

/// Project a well-formed GlobalType and compile to CodeImage.
/// Returns None if projection fails.
pub fn code_image_from_global(global: &GlobalType) -> Option<CodeImage> {
    let projected: BTreeMap<String, LocalTypeR> = telltale_theory::projection::project_all(global)
        .ok()?
        .into_iter()
        .collect();
    Some(CodeImage::from_local_types(&projected, global))
}
