use std::{
    collections::VecDeque,
    future::Future,
    io,
    sync::{Arc, Mutex as StdMutex},
    task::{Context, Poll},
    time::{Duration, Instant},
};

use futures::future::{BoxFuture, ready};
use kube::client::Body;
use tokio::{
    sync::{mpsc, oneshot},
    task::yield_now,
    time::timeout,
};
use tonic::codegen::{
    Service,
    http::{self, header::CONTENT_TYPE},
};

use super::*;
use crate::execution::Failure;
use crate::plugin::{ExecComplete, ExecFailed, exec_output};

const TEST_TIMEOUT: Duration = Duration::from_secs(2);
const NAMESPACE: &str = "supervision-tests";
const ENVIRONMENT_ID: &str = "kube-000000000000000000000000";
const RUNNING_POD: &str =
    r#"{"apiVersion":"v1","kind":"Pod","metadata":{},"status":{"phase":"Running"}}"#;
const DELETED: &str = r#"{"apiVersion":"v1","kind":"Status","status":"Success","code":200}"#;
const NOT_FOUND: &str = r#"{"apiVersion":"v1","kind":"Status","status":"Failure","reason":"NotFound","message":"pod not found","code":404}"#;
const OOM_POD: &str = r#"{
    "apiVersion":"v1","kind":"Pod","metadata":{},
    "spec":{"containers":[{"name":"job","resources":{"limits":{"memory":"1Gi"}}}]},
    "status":{"phase":"Failed","containerStatuses":[{
        "name":"job","image":"job:test","imageID":"test","ready":false,"restartCount":0,
        "state":{"terminated":{"exitCode":137,"reason":"OOMKilled","signal":9}}
    }]}
}"#;

type RecordedRequest = (http::Method, String);

struct Step {
    method: http::Method,
    suffix: &'static str,
    reply: Reply,
}

enum Reply {
    Json(u16, &'static str),
    Pending {
        entered: oneshot::Sender<()>,
        response: oneshot::Receiver<http::Response<Body>>,
    },
}

impl Step {
    fn get(body: &'static str) -> Self {
        Self {
            method: http::Method::GET,
            suffix: "",
            reply: Reply::Json(200, body),
        }
    }

    fn not_found() -> Self {
        Self {
            reply: Reply::Json(404, NOT_FOUND),
            ..Self::get(NOT_FOUND)
        }
    }

    fn delete() -> Self {
        Self {
            method: http::Method::DELETE,
            suffix: "",
            reply: Reply::Json(200, DELETED),
        }
    }
}

struct PendingResponse {
    entered: oneshot::Receiver<()>,
    response: oneshot::Sender<http::Response<Body>>,
}

fn pending_request(suffix: &'static str) -> (Step, PendingResponse) {
    let (entered, entered_rx) = oneshot::channel();
    let (response_tx, response) = oneshot::channel();
    (
        Step {
            method: http::Method::GET,
            suffix,
            reply: Reply::Pending { entered, response },
        },
        PendingResponse {
            entered: entered_rx,
            response: response_tx,
        },
    )
}

struct ScriptedTransport {
    steps: VecDeque<Step>,
    pod_path: String,
    requests: Arc<StdMutex<Vec<RecordedRequest>>>,
}

impl Service<http::Request<Body>> for ScriptedTransport {
    type Response = http::Response<Body>;
    type Error = io::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: http::Request<Body>) -> Self::Future {
        let actual = (req.method().clone(), req.uri().path().to_owned());
        self.requests.lock().unwrap().push(actual.clone());
        let Some(step) = self.steps.pop_front() else {
            return Box::pin(ready(Err(io::Error::other(format!(
                "unexpected Kubernetes request: {actual:?}"
            )))));
        };
        let expected = (step.method, format!("{}{}", self.pod_path, step.suffix));
        if actual != expected {
            return Box::pin(ready(Err(io::Error::other(format!(
                "expected Kubernetes request {expected:?}, got {actual:?}"
            )))));
        }
        match step.reply {
            Reply::Json(status, body) => Box::pin(ready(json_response(status, body))),
            Reply::Pending { entered, response } => Box::pin(async move {
                let _ = entered.send(());
                response.await.map_err(io::Error::other)
            }),
        }
    }
}

fn json_response(status: u16, body: &str) -> Result<http::Response<Body>, io::Error> {
    http::Response::builder()
        .status(status)
        .header(CONTENT_TYPE, "application/json")
        .body(Body::from(body.as_bytes().to_vec()))
        .map_err(io::Error::other)
}

struct Fixture {
    server: PluginServer,
    requests: Arc<StdMutex<Vec<RecordedRequest>>>,
    expected: Vec<RecordedRequest>,
}

impl Fixture {
    fn new(steps: Vec<Step>) -> Self {
        let pod_path = format!(
            "/api/v1/namespaces/{NAMESPACE}/pods/{}",
            JobEnvironments::pod_name(ENVIRONMENT_ID)
        );
        let expected = steps
            .iter()
            .map(|step| (step.method.clone(), format!("{pod_path}{}", step.suffix)))
            .collect();
        let requests = Arc::new(StdMutex::new(Vec::new()));
        let transport = ScriptedTransport {
            steps: steps.into(),
            pod_path,
            requests: Arc::clone(&requests),
        };
        let config = Config {
            default_job_image: "job:test".to_owned(),
            job_node_selector: None,
            job_affinity: None,
            job_tolerations: None,
            job_resource_requirements: None,
            pod_resource_requirements: None,
            service_resource_requirements: None,
            job_service_account: None,
            listen_addr: ([127, 0, 0, 1], 0).into(),
            max_environment_timeout_seconds: 60,
            namespace: NAMESPACE.to_owned(),
        };
        Self {
            server: PluginServer::new(Client::new(transport, NAMESPACE), config),
            requests,
            expected,
        }
    }

    fn assert_requests(&self, count: usize) {
        assert_eq!(*self.requests.lock().unwrap(), self.expected[..count]);
    }

    fn assert_finished(&self) {
        self.assert_requests(self.expected.len());
    }

    async fn remove(&self) {
        within(self.server.remove(Request::new(RemoveRequest {
            environment_id: ENVIRONMENT_ID.to_owned(),
        })))
        .await
        .expect("Remove must succeed after cancellation releases the environment lock");
    }
}

async fn within<F>(future: F) -> F::Output
where
    F: Future,
{
    timeout(TEST_TIMEOUT, future)
        .await
        .expect("supervision operation timed out")
}

async fn wait_for_senders<T>(sender: mpsc::WeakSender<T>) {
    // The guard owns the last sender. Waiting for its release also waits for any
    // deletion it performs, so negative DELETE assertions cannot race the guard.
    within(async {
        while sender.strong_count() != 0 {
            yield_now().await;
        }
    })
    .await;
}

fn copy_out_request() -> CopyOutRequest {
    CopyOutRequest {
        environment_id: ENVIRONMENT_ID.to_owned(),
        src_path: "/shared/output".to_owned(),
    }
}

async fn finish_output(
    server: &PluginServer,
    outcome: Result<i32, Failure>,
) -> exec_output::Output {
    let (sender, mut receiver) = mpsc::channel(1);
    let (finished, finished_rx) = oneshot::channel();
    let mut completion = Some(finished);
    within(server.workflow.finish_exec(
        ENVIRONMENT_ID,
        "test-command",
        Instant::now(),
        outcome,
        sender,
        &mut completion,
    ))
    .await
    .expect("terminal outcome should be published");
    within(finished_rx)
        .await
        .expect("completion must be explicit");
    assert!(completion.is_none());
    let output = within(receiver.recv())
        .await
        .expect("terminal message")
        .expect("terminal outcome is not a gRPC error")
        .output
        .expect("terminal output variant");
    assert!(within(receiver.recv()).await.is_none());
    output
}

#[tokio::test]
async fn ordinary_nonzero_exit_is_exec_complete() {
    let fixture = Fixture::new(vec![]);
    assert_eq!(
        finish_output(&fixture.server, Ok(23)).await,
        exec_output::Output::ExecComplete(ExecComplete { exit_code: 23 })
    );
    fixture.assert_finished();
}

#[tokio::test]
async fn remote_failure_is_exec_failed() {
    let fixture = Fixture::new(vec![]);
    let message = "Kubernetes exec ended without a status";
    assert_eq!(
        finish_output(&fixture.server, Err(Failure::Remote(message.to_owned()))).await,
        exec_output::Output::ExecFailed(ExecFailed {
            error_message: message.to_owned(),
        })
    );
    fixture.assert_finished();
}

#[tokio::test]
async fn stream_failure_remains_grpc_error_without_completion() {
    let fixture = Fixture::new(vec![]);
    for expected in [
        Status::data_loss("stdout read failed"),
        Status::cancelled("Exec stream closed"),
    ] {
        let (sender, mut receiver) = mpsc::channel(1);
        let (finished, mut finished_rx) = oneshot::channel();
        let mut completion = Some(finished);
        let error = within(fixture.server.workflow.finish_exec(
            ENVIRONMENT_ID,
            "test-command",
            Instant::now(),
            Err(Failure::Stream(expected.clone())),
            sender,
            &mut completion,
        ))
        .await
        .expect_err("stream errors must propagate rather than becoming ExecFailed");
        assert_eq!(error.code(), expected.code());
        assert_eq!(error.message(), expected.message());
        assert!(completion.is_some());
        assert_eq!(
            finished_rx.try_recv(),
            Err(oneshot::error::TryRecvError::Empty)
        );
        assert!(within(receiver.recv()).await.is_none());
    }
    fixture.assert_finished();
}

#[tokio::test]
async fn completion_precedes_terminal_publication_under_backpressure() {
    let fixture = Fixture::new(vec![]);
    for outcome in [Ok(0), Err(Failure::Remote("remote failure".to_owned()))] {
        let (sender, mut receiver) = mpsc::channel(1);
        let permit = sender.clone().try_reserve_owned().unwrap();
        let (finished, mut finished_rx) = oneshot::channel();
        let mut completion = Some(finished);
        let finishing = fixture.server.workflow.finish_exec(
            ENVIRONMENT_ID,
            "test-command",
            Instant::now(),
            outcome,
            sender,
            &mut completion,
        );
        tokio::pin!(finishing);
        assert!(futures::poll!(finishing.as_mut()).is_pending());
        assert_eq!(finished_rx.try_recv(), Ok(()));
        assert!(matches!(
            receiver.try_recv(),
            Err(mpsc::error::TryRecvError::Empty)
        ));
        drop(permit);
        within(finishing).await.unwrap();
        within(receiver.recv()).await.unwrap().unwrap();
        assert!(within(receiver.recv()).await.is_none());
    }
    fixture.assert_finished();
}

#[tokio::test]
async fn receiver_closure_during_terminal_send_does_not_cancel_completed_exec() {
    let fixture = Fixture::new(vec![]);
    for outcome in [Ok(0), Err(Failure::Remote("remote failure".to_owned()))] {
        let (sender, receiver) = mpsc::channel(1);
        let weak = sender.downgrade();
        let permit = sender.clone().try_reserve_owned().unwrap();
        let mut completion = Some(
            fixture
                .server
                .exec_cancellation_guard(ENVIRONMENT_ID.to_owned(), sender.clone()),
        );
        let finishing = fixture.server.workflow.finish_exec(
            ENVIRONMENT_ID,
            "test-command",
            Instant::now(),
            outcome,
            sender,
            &mut completion,
        );
        tokio::pin!(finishing);
        assert!(futures::poll!(finishing.as_mut()).is_pending());
        drop(receiver);
        drop(permit);
        assert_eq!(
            within(finishing).await.unwrap_err().code(),
            tonic::Code::Cancelled
        );
        wait_for_senders(weak).await;
        assert!(
            !fixture
                .server
                .environments
                .take_cancelled_environment(ENVIRONMENT_ID)
        );
    }
    fixture.assert_finished();
}

#[tokio::test]
async fn explicit_completion_wins_when_receiver_closure_is_also_ready() {
    let fixture = Fixture::new(vec![]);
    for _ in 0..32 {
        let (sender, receiver) = mpsc::channel::<Result<ExecOutput, Status>>(1);
        let weak = sender.downgrade();
        let finished = fixture
            .server
            .exec_cancellation_guard(ENVIRONMENT_ID.to_owned(), sender);
        // Neither operation yields: both guard branches are ready on its first poll.
        finished.send(()).unwrap();
        drop(receiver);
        wait_for_senders(weak).await;
        assert!(
            !fixture
                .server
                .environments
                .take_cancelled_environment(ENVIRONMENT_ID)
        );
    }
    fixture.assert_finished();
}

#[tokio::test]
async fn exit_137_uses_pod_termination_diagnosis() {
    let fixture = Fixture::new(vec![Step::get(OOM_POD)]);
    let output = finish_output(&fixture.server, Ok(137)).await;
    let message = match output {
        exec_output::Output::ExecFailed(failure) => failure.error_message,
        other => panic!("expected ExecFailed, got {other:?}"),
    };
    assert!(message.contains("reason=OOMKilled"));
    assert!(message.contains("memory_limit=1Gi"));
    assert!(message.contains("increase KURANTO_JOB_MEMORY_LIMIT"));
    fixture.assert_finished();
}

#[tokio::test]
async fn stopped_container_remote_failure_uses_pod_termination_diagnosis() {
    let fixture = Fixture::new(vec![Step::get(OOM_POD)]);
    let output = finish_output(
        &fixture.server,
        Err(Failure::Remote("cannot exec in a stopped state".to_owned())),
    )
    .await;
    let message = match output {
        exec_output::Output::ExecFailed(failure) => failure.error_message,
        other => panic!("expected ExecFailed, got {other:?}"),
    };
    assert!(message.contains("reason=OOMKilled"));
    assert!(message.contains("exit_code=137"));
    fixture.assert_finished();
}

#[tokio::test]
async fn copy_out_cancellation_while_waiting_for_lock_does_not_delete_pod() {
    let fixture = Fixture::new(vec![
        Step::get(RUNNING_POD),
        Step::delete(),
        Step::not_found(),
    ]);
    let held = within(fixture.server.environments.acquire(ENVIRONMENT_ID))
        .await
        .expect("environment lock should be acquired");
    let (sender, receiver) = mpsc::channel(1);
    let weak = sender.downgrade();
    let copying = fixture
        .server
        .clone()
        .workflow
        .run_copy_out(copy_out_request(), sender);
    tokio::pin!(copying);
    assert!(futures::poll!(copying.as_mut()).is_pending());
    drop(receiver);
    within(copying).await;
    wait_for_senders(weak).await;
    fixture.assert_requests(0);
    assert!(
        !fixture
            .server
            .environments
            .take_cancelled_environment(ENVIRONMENT_ID)
    );

    let removing = fixture.remove();
    tokio::pin!(removing);
    assert!(futures::poll!(removing.as_mut()).is_pending());
    fixture.assert_requests(0);
    drop(held);
    within(removing).await;
    fixture.assert_finished();
}

#[tokio::test]
async fn copy_out_cancellation_during_attach_releases_lock_without_deleting_pod() {
    let (attach, mut pending) = pending_request("/exec");
    let fixture = Fixture::new(vec![
        Step::get(RUNNING_POD),
        attach,
        Step::get(RUNNING_POD),
        Step::delete(),
        Step::not_found(),
    ]);
    let (sender, receiver) = mpsc::channel(1);
    let weak = sender.downgrade();
    let copying = tokio::spawn(
        fixture
            .server
            .clone()
            .workflow
            .run_copy_out(copy_out_request(), sender),
    );
    within(&mut pending.entered).await.unwrap();
    drop(receiver);
    within(copying).await.unwrap();
    within(pending.response.closed()).await;
    wait_for_senders(weak).await;
    fixture.assert_requests(2);
    assert!(
        !fixture
            .server
            .environments
            .take_cancelled_environment(ENVIRONMENT_ID)
    );
    fixture.remove().await;
    fixture.assert_finished();
}

#[tokio::test]
async fn silent_exec_cancellation_during_attach_still_deletes_pod() {
    let (attach, mut pending) = pending_request("/exec");
    let (deletion_check, mut deletion) = pending_request("");
    let fixture = Fixture::new(vec![
        Step::get(RUNNING_POD),
        attach,
        Step::delete(),
        deletion_check,
        Step::not_found(),
    ]);
    let (sender, mut receiver) = mpsc::channel(1);
    let weak = sender.downgrade();
    let executing = tokio::spawn(fixture.server.clone().run_exec(
        ExecRequest {
            environment_id: ENVIRONMENT_ID.to_owned(),
            command: vec!["silent-command".to_owned()],
            ..Default::default()
        },
        sender,
    ));
    within(&mut pending.entered).await.unwrap();
    assert!(matches!(
        receiver.try_recv(),
        Err(mpsc::error::TryRecvError::Empty)
    ));
    drop(receiver);
    // Observe DELETE before releasing attach: cancellation cannot depend on a
    // failed output write, an attach response, or acquiring the environment lock.
    within(&mut deletion.entered).await.unwrap();
    fixture.assert_requests(4);
    deletion
        .response
        .send(json_response(404, NOT_FOUND).unwrap())
        .unwrap();
    drop(pending.response);
    within(executing).await.unwrap();
    wait_for_senders(weak).await;
    fixture.remove().await;
    fixture.assert_finished();
}

#[tokio::test]
async fn dropped_completion_sender_does_not_mask_exec_cancellation() {
    let fixture = Fixture::new(vec![Step::delete(), Step::not_found()]);
    let (sender, receiver) = mpsc::channel::<Result<ExecOutput, Status>>(1);
    let weak = sender.downgrade();
    let finished = fixture
        .server
        .exec_cancellation_guard(ENVIRONMENT_ID.to_owned(), sender);
    // Dropping the completion sender is not an explicit completion notification.
    // Both branches are ready before the guard gets its first poll.
    drop(finished);
    drop(receiver);
    wait_for_senders(weak).await;
    assert!(
        fixture
            .server
            .environments
            .take_cancelled_environment(ENVIRONMENT_ID)
    );
    fixture.assert_finished();
}