use super::*;
use crate::model::ContentBlock;
use serde_json::json;
use std::sync::{Arc, Mutex as StdMutex};
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;

fn body(input: Vec<Value>) -> Value {
    json!({
        "model": "gpt-5-codex",
        "instructions": "system",
        "input": input,
        "store": false,
        "stream": true,
        "prompt_cache_key": "rho:session",
        "tools": [{"type":"function","name":"read","parameters":{"type":"object"}}],
        "tool_choice": "auto",
        "reasoning": {"effort":"low","summary":"auto"},
    })
}

fn candidate(input: Vec<Value>) -> CodexContinuationCandidate {
    CodexContinuationCandidate::from_responses_body(&body(input)).unwrap()
}

fn tokens() -> CodexTokens {
    CodexTokens {
        access_token: "token".into(),
        refresh_token: None,
        id_token: None,
        account_id: Some("account".into()),
    }
}

async fn ws_server(expected_messages: usize) -> (String, Arc<StdMutex<Vec<Value>>>) {
    ws_server_connections(vec![expected_messages]).await
}

async fn ws_server_connections(
    expected_messages_by_connection: Vec<usize>,
) -> (String, Arc<StdMutex<Vec<Value>>>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let frames = Arc::new(StdMutex::new(Vec::new()));
    let server_frames = Arc::clone(&frames);
    tokio::spawn(async move {
        let mut response_index = 0;
        for expected_messages in expected_messages_by_connection {
            let (stream, _) = listener.accept().await.unwrap();
            let mut socket = accept_async(stream).await.unwrap();
            for _ in 0..expected_messages {
                response_index += 1;
                let message = socket.next().await.unwrap().unwrap();
                let text = message.into_text().unwrap();
                let frame: Value = serde_json::from_str(&text).unwrap();
                server_frames.lock().unwrap().push(frame);
                let response_id = format!("resp_{response_index}");
                socket
                    .send(Message::Text(
                        json!({"type":"response.output_text.delta","delta":format!("ok{response_index}")})
                            .to_string()
                            .into(),
                    ))
                    .await
                    .unwrap();
                socket
                    .send(Message::Text(
                        json!({
                            "type":"response.completed",
                            "response":{
                                "id": response_id,
                                "output_text": format!("ok{response_index}"),
                                "output":[],
                                "usage":{"input_tokens": 10, "output_tokens": 2}
                            }
                        })
                        .to_string()
                        .into(),
                    ))
                    .await
                    .unwrap();
            }
        }
    });
    (format!("ws://{addr}/responses"), frames)
}

async fn ws_server_closes_after_delta() -> (String, Arc<StdMutex<Vec<Value>>>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let frames = Arc::new(StdMutex::new(Vec::new()));
    let server_frames = Arc::clone(&frames);
    tokio::spawn(async move {
        let (stream, _) = listener.accept().await.unwrap();
        let mut socket = accept_async(stream).await.unwrap();
        let message = socket.next().await.unwrap().unwrap();
        let text = message.into_text().unwrap();
        let frame: Value = serde_json::from_str(&text).unwrap();
        server_frames.lock().unwrap().push(frame);
        socket
            .send(Message::Text(
                json!({"type":"response.output_text.delta","delta":"partial"})
                    .to_string()
                    .into(),
            ))
            .await
            .unwrap();
    });
    (format!("ws://{addr}/responses"), frames)
}

async fn ws_server_stalls_after_request() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        let (stream, _) = listener.accept().await.unwrap();
        let mut socket = accept_async(stream).await.unwrap();
        let _request = socket.next().await.unwrap().unwrap();
        std::future::pending::<()>().await;
    });
    format!("ws://{addr}/responses")
}

async fn ws_server_sends_keep_alive_frames() -> String {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        let (stream, _) = listener.accept().await.unwrap();
        let mut socket = accept_async(stream).await.unwrap();
        let _request = socket.next().await.unwrap().unwrap();
        loop {
            socket.send(Message::Ping(Vec::new().into())).await.unwrap();
            tokio::time::sleep(std::time::Duration::from_millis(1)).await;
        }
    });
    format!("ws://{addr}/responses")
}

#[test]
fn builds_delta_body_when_next_input_extends_previous_input() {
    let first = candidate(vec![json!({"role":"user","content":"one"})]);
    let second = candidate(vec![
        json!({"role":"user","content":"one"}),
        json!({"role":"assistant","content":"two"}),
        json!({"role":"user","content":"three"}),
    ]);
    let mut state = CodexContinuationState::default();
    state.record_success(&first, Some("resp_1".into()));

    let plan = state.plan_delta(&second);

    let CodexContinuationPlan::Delta {
        previous_response_id,
        input,
        body,
    } = plan
    else {
        panic!("expected delta plan");
    };
    assert_eq!(previous_response_id, "resp_1");
    assert_eq!(
        input,
        vec![
            json!({"role":"assistant","content":"two"}),
            json!({"role":"user","content":"three"}),
        ]
    );
    assert_eq!(body["previous_response_id"], "resp_1");
    assert_eq!(body["input"], json!(input));
    assert_eq!(body["model"], "gpt-5-codex");
    assert_eq!(body["prompt_cache_key"], "rho:session");
    assert_eq!(body["tools"][0]["name"], "read");
    assert_eq!(body["reasoning"], json!({"effort":"low","summary":"auto"}));
    assert_eq!(body["store"], false);
    assert_eq!(body["stream"], true);
}

#[tokio::test]
async fn responses_lite_websocket_request_sets_lite_client_metadata() {
    let (url, frames) = ws_server(1).await;
    let transport = CodexWsTransport::new_with_url(url);
    let mut on_event = None;

    transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::ResponsesLite,
            &mut on_event,
        )
        .await
        .unwrap();

    let frames = frames.lock().unwrap();
    assert_eq!(
        frames[0]["client_metadata"]["ws_request_header_x_openai_internal_codex_responses_lite"],
        "true"
    );
}

#[tokio::test]
async fn responses_lite_websocket_requests_do_not_use_incomplete_continuation_state() {
    let (url, frames) = ws_server(2).await;
    let transport = CodexWsTransport::new_with_url(url);
    let mut on_event = None;

    transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::ResponsesLite,
            &mut on_event,
        )
        .await
        .unwrap();
    transport
        .send_responses_turn(
            body(vec![
                json!({"role":"user","content":"one"}),
                json!({"role":"assistant","content":"two"}),
                json!({"role":"user","content":"three"}),
            ]),
            &tokens(),
            CodexRequestMode::ResponsesLite,
            &mut on_event,
        )
        .await
        .unwrap();

    let frames = frames.lock().unwrap();
    assert_eq!(frames.len(), 2);
    assert!(frames[1].get("previous_response_id").is_none());
    assert_eq!(
        frames[1]["input"],
        json!([
            {"role":"user","content":"one"},
            {"role":"assistant","content":"two"},
            {"role":"user","content":"three"}
        ])
    );
}

#[tokio::test]
async fn first_websocket_request_sends_full_input_without_previous_response_id() {
    let (url, frames) = ws_server(1).await;
    let transport = CodexWsTransport::new_with_url(url);
    let mut on_event = None;

    let turn = transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();

    assert!(matches!(turn, CodexWsTurn::Completed(_)));
    let frames = frames.lock().unwrap();
    assert_eq!(frames.len(), 1);
    assert_eq!(frames[0]["type"], "response.create");
    assert!(frames[0].get("previous_response_id").is_none());
    assert_eq!(frames[0]["input"], json!([{"role":"user","content":"one"}]));
}

#[tokio::test]
async fn compatible_websocket_request_sends_delta_with_previous_response_id() {
    let (url, frames) = ws_server(2).await;
    let transport = CodexWsTransport::new_with_url(url);
    let mut on_event = None;

    transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();
    let turn = transport
        .send_responses_turn(
            body(vec![
                json!({"role":"user","content":"one"}),
                json!({"role":"assistant","content":"two"}),
                json!({"role":"user","content":"three"}),
            ]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();

    let CodexWsTurn::Completed(ModelResponse::Assistant(blocks)) = turn else {
        panic!("expected websocket completion");
    };
    assert!(matches!(
        blocks.as_slice(),
        [ContentBlock::Text(text)] if text == "ok2"
    ));
    let frames = frames.lock().unwrap();
    assert_eq!(frames.len(), 2);
    assert_eq!(frames[1]["previous_response_id"], "resp_1");
    assert_eq!(
        frames[1]["input"],
        json!([
            {"role":"assistant","content":"two"},
            {"role":"user","content":"three"}
        ])
    );
}

#[test]
fn falls_back_to_full_request_without_previous_response_id() {
    let state = CodexContinuationState::default();
    let plan = state.plan_delta(&candidate(vec![json!({"role":"user","content":"one"})]));

    assert_eq!(
        plan,
        CodexContinuationPlan::Full {
            reason: CodexContinuationFullReason::MissingPreviousResponse
        }
    );
}

#[test]
fn resets_when_history_is_rewritten_by_compaction() {
    let first = candidate(vec![
        json!({"role":"user","content":"old"}),
        json!({"role":"assistant","content":"answer"}),
    ]);
    let compacted_body = body(vec![
        json!({"role":"user","content":"summary of old conversation"}),
        json!({"role":"user","content":"new"}),
    ]);
    let compacted = CodexContinuationCandidate::from_responses_body(&compacted_body).unwrap();
    let mut state = CodexContinuationState::default();
    state.record_success(&first, Some("resp_1".into()));

    let plan = state.plan_request(&compacted, compacted_body.clone());

    assert!(!plan.planned_delta);
    assert_eq!(
        plan.reset_reason,
        Some(CodexContinuationResetReason::HistoryRewritten)
    );
    assert_eq!(plan.body, compacted_body);
    assert_eq!(
        state.plan_delta(&compacted),
        CodexContinuationPlan::Full {
            reason: CodexContinuationFullReason::MissingPreviousResponse
        }
    );
}

#[test]
fn resets_when_tools_change() {
    let first = candidate(vec![json!({"role":"user","content":"one"})]);
    let mut changed_body = body(vec![
        json!({"role":"user","content":"one"}),
        json!({"role":"user","content":"two"}),
    ]);
    changed_body["tools"] = json!([{ "type":"function", "name":"write" }]);
    let changed = CodexContinuationCandidate::from_responses_body(&changed_body).unwrap();
    let mut state = CodexContinuationState::default();
    state.record_success(&first, Some("resp_1".into()));

    let plan = state.plan_request(&changed, changed_body);

    assert_eq!(
        plan.reset_reason,
        Some(CodexContinuationResetReason::ToolsChanged)
    );
    assert_eq!(
        state.plan_delta(&changed),
        CodexContinuationPlan::Full {
            reason: CodexContinuationFullReason::MissingPreviousResponse
        }
    );
}

#[test]
fn resets_when_model_changes() {
    assert_reset_reason_for_body_change(
        |body| body["model"] = json!("gpt-5-codex-alt"),
        CodexContinuationResetReason::ModelChanged,
    );
}

#[test]
fn resets_when_reasoning_changes() {
    assert_reset_reason_for_body_change(
        |body| body["reasoning"] = json!({"effort":"high","summary":"auto"}),
        CodexContinuationResetReason::ReasoningChanged,
    );
}

#[test]
fn resets_when_prompt_cache_key_changes() {
    assert_reset_reason_for_body_change(
        |body| body["prompt_cache_key"] = json!("rho:other"),
        CodexContinuationResetReason::PromptCacheKeyChanged,
    );
}

#[test]
fn resets_when_tool_choice_changes() {
    assert_reset_reason_for_body_change(
        |body| body["tool_choice"] = json!("none"),
        CodexContinuationResetReason::ToolChoiceChanged,
    );
}

fn assert_reset_reason_for_body_change(
    mutate: impl FnOnce(&mut Value),
    expected: CodexContinuationResetReason,
) {
    let first = candidate(vec![json!({"role":"user","content":"one"})]);
    let mut changed_body = body(vec![
        json!({"role":"user","content":"one"}),
        json!({"role":"user","content":"two"}),
    ]);
    mutate(&mut changed_body);
    let changed = CodexContinuationCandidate::from_responses_body(&changed_body).unwrap();
    let mut state = CodexContinuationState::default();
    state.record_success(&first, Some("resp_1".into()));

    let plan = state.plan_request(&changed, changed_body);

    assert_eq!(plan.reset_reason, Some(expected));
}

#[tokio::test]
async fn stalled_websocket_falls_back_instead_of_blocking_the_turn() {
    let url = ws_server_stalls_after_request().await;
    let mut transport = CodexWsTransport::new_with_url(url);
    transport.idle_timeout = std::time::Duration::from_millis(10);
    let mut on_event = None;

    let outcome = transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();

    assert!(matches!(outcome, CodexWsTurn::FullSseFallback));
}

#[tokio::test]
async fn websocket_keep_alive_frames_do_not_reset_the_idle_timeout() {
    let url = ws_server_sends_keep_alive_frames().await;
    let mut transport = CodexWsTransport::new_with_url(url);
    transport.idle_timeout = std::time::Duration::from_millis(10);
    let mut on_event = None;

    let outcome = transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();

    assert!(matches!(outcome, CodexWsTurn::FullSseFallback));
}

#[tokio::test]
async fn websocket_error_resets_continuation_and_returns_full_sse_fallback() {
    let (url, frames) = ws_server_connections(vec![1, 1]).await;
    let transport = CodexWsTransport::new_with_url(url);
    let mut on_event = None;
    transport
        .send_responses_turn(
            body(vec![json!({"role":"user","content":"one"})]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();

    let outcome = transport
        .send_responses_turn(
            body(vec![
                json!({"role":"user","content":"one"}),
                json!({"role":"user","content":"two"}),
            ]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();

    assert!(matches!(outcome, CodexWsTurn::FullSseFallback));

    transport
        .send_responses_turn(
            body(vec![
                json!({"role":"user","content":"one"}),
                json!({"role":"user","content":"two"}),
            ]),
            &tokens(),
            CodexRequestMode::Standard,
            &mut on_event,
        )
        .await
        .unwrap();
    let frames = frames.lock().unwrap();
    assert_eq!(frames.len(), 2);
    assert!(frames[1].get("previous_response_id").is_none());
}

#[tokio::test]
async fn websocket_fallback_does_not_emit_partial_events() {
    let (url, frames) = ws_server_closes_after_delta().await;
    let transport = CodexWsTransport::new_with_url(url);
    let mut deltas = Vec::new();
    {
        let mut collect_event = |event| {
            if let ModelEvent::OutputDelta(delta) = event {
                deltas.push(delta);
            }
            Ok(())
        };
        let mut on_event =
            Some(&mut collect_event as &mut dyn FnMut(ModelEvent) -> Result<(), ModelError>);

        let outcome = transport
            .send_responses_turn(
                body(vec![json!({"role":"user","content":"one"})]),
                &tokens(),
                CodexRequestMode::Standard,
                &mut on_event,
            )
            .await
            .unwrap();

        assert!(matches!(outcome, CodexWsTurn::FullSseFallback));
    }

    assert!(deltas.is_empty());
    assert_eq!(frames.lock().unwrap().len(), 1);
}

#[test]
fn derives_websocket_url_from_codex_api_base() {
    assert_eq!(
        codex_ws_url("https://chatgpt.com/backend-api/codex"),
        "wss://chatgpt.com/backend-api/codex/responses"
    );
    assert_eq!(
        codex_ws_url("http://127.0.0.1:1234/codex/"),
        "ws://127.0.0.1:1234/codex/responses"
    );
}
