use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};

use async_trait::async_trait;
use motosan_agent_loop::{
    AgentEvent, AgentLoop, AgentOp, AgentSession, BackpressurePolicy, ChannelConfig, ChatOutput,
    LlmClient, LlmResponse, Message,
};
use motosan_agent_tool::{Tool, ToolContext, ToolDef, ToolResult};

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// An LLM that runs multiple tool-call iterations before returning a final
/// message.  Each tool call introduces a short delay, simulating real latency.
struct SlowIteratingLlm {
    tool_iterations: usize,
    call_count: AtomicUsize,
}

impl SlowIteratingLlm {
    fn new(tool_iterations: usize) -> Self {
        Self {
            tool_iterations,
            call_count: AtomicUsize::new(0),
        }
    }
}

#[async_trait]
impl LlmClient for SlowIteratingLlm {
    async fn chat(
        &self,
        _messages: &[Message],
        _tools: &[ToolDef],
    ) -> motosan_agent_loop::Result<ChatOutput> {
        let n = self.call_count.fetch_add(1, Ordering::SeqCst);
        if n < self.tool_iterations {
            Ok(ChatOutput::new(LlmResponse::single_tool_call(
                format!("call_{n}"),
                "slow_tool".to_string(),
                serde_json::json!({}),
            )))
        } else {
            Ok(ChatOutput::new(LlmResponse::Message("done".to_string())))
        }
    }
}

struct SlowTool {
    delay_ms: u64,
}

impl Tool for SlowTool {
    fn def(&self) -> ToolDef {
        ToolDef {
            name: "slow_tool".to_string(),
            description: "slow".to_string(),
            input_schema: serde_json::json!({
                "type": "object",
                "properties": {},
                "required": []
            }),
        }
    }

    fn call(
        &self,
        _args: serde_json::Value,
        _ctx: &ToolContext,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ToolResult> + Send + '_>> {
        let delay = self.delay_ms;
        Box::pin(async move {
            tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
            ToolResult::text("ok")
        })
    }
}

/// A simple mock LLM that returns pre-programmed responses.
struct MockLlm {
    responses: Mutex<Vec<LlmResponse>>,
}

impl MockLlm {
    fn new(responses: Vec<LlmResponse>) -> Self {
        Self {
            responses: Mutex::new(responses),
        }
    }
}

#[async_trait]
impl LlmClient for MockLlm {
    async fn chat(
        &self,
        _messages: &[Message],
        _tools: &[ToolDef],
    ) -> motosan_agent_loop::Result<ChatOutput> {
        let mut responses = self.responses.lock().unwrap();
        if responses.is_empty() {
            panic!("MockLlm: no more responses");
        }
        Ok(ChatOutput::new(responses.remove(0)))
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

/// Verify that channel capacities are configurable via the builder and that
/// sessions created with custom config use those capacities.
#[tokio::test]
async fn configurable_channel_capacities() {
    let config = ChannelConfig {
        input_capacity: 4,
        ops_capacity: 2,
        ops_backpressure: BackpressurePolicy::Block,
    };

    let agent = AgentLoop::builder()
        .channel_config(config.clone())
        .max_iterations(5)
        .build();

    assert_eq!(agent.channel_config().input_capacity, 4);
    assert_eq!(agent.channel_config().ops_capacity, 2);
    assert_eq!(
        agent.channel_config().ops_backpressure,
        BackpressurePolicy::Block
    );
}

/// Verify that individual builder methods for channel config work.
#[tokio::test]
async fn builder_individual_channel_methods() {
    let agent = AgentLoop::builder()
        .input_channel_capacity(8)
        .ops_channel_capacity(4)
        .ops_backpressure(BackpressurePolicy::Reject)
        .build();

    assert_eq!(agent.channel_config().input_capacity, 8);
    assert_eq!(agent.channel_config().ops_capacity, 4);
    assert_eq!(
        agent.channel_config().ops_backpressure,
        BackpressurePolicy::Reject
    );
}

/// Default channel config preserves backward-compatible values.
#[tokio::test]
async fn default_channel_config_backward_compatible() {
    let config = ChannelConfig::default();
    assert_eq!(config.input_capacity, 64);
    assert_eq!(config.ops_capacity, 64);
    assert_eq!(config.ops_backpressure, BackpressurePolicy::Block);
}

/// Stress test: with a tiny bounded ops channel (capacity=2) and Reject policy,
/// rapid fire ops that exceed capacity are rejected rather than blocking.
/// Memory should remain bounded since the channel never grows.
/// Also verifies that OpRejected and OpsSaturated telemetry events are emitted.
#[tokio::test]
async fn reject_policy_does_not_block_under_overload() {
    let llm = Arc::new(SlowIteratingLlm::new(3));
    let tool: Arc<dyn Tool> = Arc::new(SlowTool { delay_ms: 50 });

    let agent_loop = AgentLoop::builder()
        .tool(tool)
        .max_iterations(5)
        .ops_channel_capacity(2)
        .ops_backpressure(BackpressurePolicy::Reject)
        .build();

    let rejected_count = Arc::new(AtomicUsize::new(0));
    let rejected_clone = rejected_count.clone();
    let op_rejected_events = Arc::new(AtomicUsize::new(0));
    let ops_saturated_events = Arc::new(AtomicUsize::new(0));
    let op_rejected_clone = op_rejected_events.clone();
    let ops_saturated_clone = ops_saturated_events.clone();

    let (session, handle) = AgentSession::new(agent_loop, llm, move |event| match &event {
        AgentEvent::OpRejected { .. } => {
            op_rejected_clone.fetch_add(1, Ordering::SeqCst);
        }
        AgentEvent::OpsSaturated { .. } => {
            ops_saturated_clone.fetch_add(1, Ordering::SeqCst);
        }
        _ => {}
    });

    session.send("start").await;

    // Fire many ops rapidly -- most should be rejected since the channel is
    // tiny and the loop is busy with tool calls.
    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
    let mut sent = 0;
    let mut rejected = 0;
    for i in 0..100 {
        let ok = session
            .send_op(AgentOp::InjectHint(format!("hint-{i}")))
            .await;
        sent += 1;
        if !ok {
            rejected += 1;
        }
    }
    rejected_clone.store(rejected, Ordering::SeqCst);

    // Wait for the session to finish processing.
    drop(session);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(5), handle).await;

    // Under Reject policy, at least some ops should have been rejected
    // (the channel capacity is 2 and we sent 100 rapidly).
    let total_rejected = rejected_count.load(Ordering::SeqCst);
    assert!(
        total_rejected > 0,
        "expected some ops to be rejected, but none were (sent={sent})"
    );

    // Verify telemetry events were emitted for rejected ops.
    let total_op_rejected = op_rejected_events.load(Ordering::SeqCst);
    assert!(
        total_op_rejected > 0,
        "expected OpRejected events to be emitted, but got none"
    );
    assert_eq!(
        total_op_rejected, total_rejected,
        "OpRejected event count should match rejected op count"
    );

    // OpsSaturated should fire exactly once.
    let total_ops_saturated = ops_saturated_events.load(Ordering::SeqCst);
    assert_eq!(
        total_ops_saturated, 1,
        "expected exactly 1 OpsSaturated event, got {total_ops_saturated}"
    );
}

/// Stress test: with Block policy and small channel, producer-consumer
/// imbalance does not cause unbounded memory growth. The test verifies that
/// the session completes without hanging (timeout enforced).
#[tokio::test]
async fn block_policy_completes_under_load() {
    let llm = Arc::new(SlowIteratingLlm::new(2));
    let tool: Arc<dyn Tool> = Arc::new(SlowTool { delay_ms: 10 });

    let agent_loop = AgentLoop::builder()
        .tool(tool)
        .max_iterations(5)
        .ops_channel_capacity(2)
        .ops_backpressure(BackpressurePolicy::Block)
        .build();

    let events_received = Arc::new(AtomicUsize::new(0));
    let events_clone = events_received.clone();

    let (session, handle) = AgentSession::new(agent_loop, llm, move |_event| {
        events_clone.fetch_add(1, Ordering::SeqCst);
    });

    session.send("start").await;

    // Give the background loop time to pick up the message before closing.
    // `drop` sets the closed flag which causes the loop to skip unprocessed
    // messages, so we need to ensure processing has started.
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // Wait for the session to finish.
    drop(session);
    let result = tokio::time::timeout(std::time::Duration::from_secs(10), handle).await;
    assert!(
        result.is_ok(),
        "session should complete within timeout under Block policy"
    );

    // Verify events were actually emitted.
    assert!(
        events_received.load(Ordering::SeqCst) > 0,
        "expected at least some events to be emitted"
    );
}

/// Verify that the new AgentEvent variants for telemetry are constructable
/// and have the expected Debug output.
#[test]
fn telemetry_event_variants_are_constructable() {
    let events: Vec<AgentEvent> = vec![
        AgentEvent::OpsSaturated { capacity: 64 },
        AgentEvent::OpDropped {
            reason: "channel full".to_string(),
        },
        AgentEvent::OpRejected {
            reason: "backpressure".to_string(),
        },
    ];

    for event in &events {
        // Verify Debug is implemented and doesn't panic.
        let _ = format!("{event:?}");
    }

    // Verify Clone is implemented.
    let cloned = events[0].clone();
    let _ = format!("{cloned:?}");
}

/// Stress test: many concurrent producers sending to a bounded session.
/// With Reject policy, all producers should return quickly (no deadlock).
#[tokio::test]
async fn concurrent_producers_reject_policy_no_deadlock() {
    let llm = Arc::new(SlowIteratingLlm::new(1));
    let tool: Arc<dyn Tool> = Arc::new(SlowTool { delay_ms: 100 });

    let agent_loop = AgentLoop::builder()
        .tool(tool)
        .max_iterations(3)
        .ops_channel_capacity(4)
        .ops_backpressure(BackpressurePolicy::Reject)
        .build();

    let (session, handle) = AgentSession::new(agent_loop, llm, |_| {});

    session.send("go").await;

    // Spawn many concurrent producers.
    let mut producer_handles = Vec::new();
    for i in 0..10 {
        let s = session.clone();
        producer_handles.push(tokio::spawn(async move {
            let mut rejected = 0usize;
            for j in 0..20 {
                if !s.send_op(AgentOp::InjectHint(format!("p{i}-h{j}"))).await {
                    rejected += 1;
                }
            }
            rejected
        }));
    }

    // All producers should finish quickly (no deadlock).
    let results = tokio::time::timeout(std::time::Duration::from_secs(5), async {
        let mut total_rejected = 0usize;
        for h in producer_handles {
            total_rejected += h.await.unwrap();
        }
        total_rejected
    })
    .await;

    assert!(results.is_ok(), "producers should complete within timeout");

    drop(session);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(5), handle).await;
}

/// Verify that BackpressurePolicy defaults match Block.
#[test]
fn backpressure_policy_default_is_block() {
    assert_eq!(BackpressurePolicy::default(), BackpressurePolicy::Block);
}

/// Verify that DropOldest policy does not block under overload.
/// Since tokio mpsc does not support true drop-oldest, it degrades to
/// drop-newest (reject) behavior, but crucially it must not block.
/// Also verifies that OpDropped and OpsSaturated telemetry events are emitted.
#[tokio::test]
async fn drop_oldest_policy_does_not_block() {
    let llm = Arc::new(SlowIteratingLlm::new(2));
    let tool: Arc<dyn Tool> = Arc::new(SlowTool { delay_ms: 50 });

    let agent_loop = AgentLoop::builder()
        .tool(tool)
        .max_iterations(5)
        .ops_channel_capacity(2)
        .ops_backpressure(BackpressurePolicy::DropOldest)
        .build();

    let op_dropped_events = Arc::new(AtomicUsize::new(0));
    let ops_saturated_events = Arc::new(AtomicUsize::new(0));
    let op_dropped_clone = op_dropped_events.clone();
    let ops_saturated_clone = ops_saturated_events.clone();

    let (session, handle) = AgentSession::new(agent_loop, llm, move |event| match &event {
        AgentEvent::OpDropped { .. } => {
            op_dropped_clone.fetch_add(1, Ordering::SeqCst);
        }
        AgentEvent::OpsSaturated { .. } => {
            ops_saturated_clone.fetch_add(1, Ordering::SeqCst);
        }
        _ => {}
    });

    session.send("start").await;
    tokio::time::sleep(std::time::Duration::from_millis(10)).await;

    // Rapid fire ops -- should not block.
    let start = std::time::Instant::now();
    let mut dropped = 0;
    for i in 0..50 {
        if !session
            .send_op(AgentOp::InjectHint(format!("hint-{i}")))
            .await
        {
            dropped += 1;
        }
    }
    let elapsed = start.elapsed();

    // With a non-blocking policy, 50 try_sends should complete in well
    // under 1 second (each is essentially instant).
    assert!(
        elapsed < std::time::Duration::from_secs(1),
        "DropOldest should not block; elapsed={elapsed:?}"
    );

    drop(session);
    let _ = tokio::time::timeout(std::time::Duration::from_secs(5), handle).await;

    // Verify OpDropped events match the number of dropped ops.
    let total_dropped_events = op_dropped_events.load(Ordering::SeqCst);
    assert_eq!(
        total_dropped_events, dropped,
        "OpDropped event count should match dropped op count"
    );

    // If any ops were dropped, OpsSaturated should fire exactly once.
    if dropped > 0 {
        let total_ops_saturated = ops_saturated_events.load(Ordering::SeqCst);
        assert_eq!(
            total_ops_saturated, 1,
            "expected exactly 1 OpsSaturated event, got {total_ops_saturated}"
        );
    }
}

/// run_with_ops still works correctly with custom channel config (backward compat).
#[tokio::test]
async fn run_with_ops_respects_custom_ops_capacity() {
    let llm = MockLlm::new(vec![LlmResponse::Message("ok".to_string())]);

    let agent = AgentLoop::builder()
        .ops_channel_capacity(8)
        .max_iterations(3)
        .build();

    // run_with_ops takes an external Receiver, so it works regardless of
    // the internal channel_config. This test ensures it still works.
    let (tx, rx) = tokio::sync::mpsc::channel::<AgentOp>(8);
    drop(tx);

    let result = agent
        .run_with_ops(&llm, vec![Message::user("hi")], Some(rx), |_| {})
        .await
        .unwrap();

    assert_eq!(result.answer, "ok");
}

/// Verify memory stays bounded: with a tiny channel, accumulating many ops
/// in a Vec would grow unboundedly, but bounded channels prevent this.
/// We check that the channel's capacity constraint is respected.
#[tokio::test]
async fn bounded_channel_limits_memory() {
    // Create a channel with capacity 4.
    let (tx, mut rx) = tokio::sync::mpsc::channel::<AgentOp>(4);

    // Fill the channel to capacity.
    for i in 0..4 {
        tx.try_send(AgentOp::InjectHint(format!("msg-{i}")))
            .unwrap();
    }

    // The 5th send should fail (channel full).
    let result = tx.try_send(AgentOp::InjectHint("overflow".to_string()));
    assert!(result.is_err(), "channel should be full at capacity 4");

    // Drain one and verify we can send again.
    let _ = rx.recv().await;
    tx.try_send(AgentOp::InjectHint("after-drain".to_string()))
        .unwrap();
}
