use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use async_trait::async_trait;
use rlmesh_proto::model::v1::{
    ConfigureRouteRequest, JoinRequest, PredictRequest, join_request, join_response,
};
use tokio::net::TcpListener;
use tokio::sync::{Mutex, mpsc};
use tokio_stream::wrappers::TcpListenerStream;

use super::lifecycle::update_lifecycle;
use super::server::{ModelRouteConfig, handle_model_request};
use super::*;
use crate::{BindAddress, ConnectAddress, Result, ServeOptions, spaces};

#[derive(Default)]
struct RecordingHandler {
    resets: Vec<(String, i32)>,
    episode_ends: Vec<ModelEpisodeEnd>,
    predictions: Vec<(String, i64, i32)>,
}

#[async_trait]
impl ModelHandler for RecordingHandler {
    async fn predict(&mut self, observation: ModelObservation) -> Result<spaces::BinaryPayload> {
        self.predictions.push((
            observation.episode_id().to_string(),
            observation.step(),
            observation.env_index(),
        ));
        Ok(spaces::BinaryPayload {
            data: vec![1, 2, 3],
        })
    }

    async fn on_reset(&mut self, observation: &ModelObservation) -> Result<()> {
        self.resets.push((
            observation.episode_id().to_string(),
            observation.env_index(),
        ));
        Ok(())
    }

    async fn on_episode_end(&mut self, event: ModelEpisodeEnd) -> Result<()> {
        self.episode_ends.push(event);
        Ok(())
    }
}

fn raw_observation(
    episode_id: &str,
    step: i64,
    env_index: i32,
    is_reset: bool,
) -> ModelObservation {
    ModelObservation {
        observation: Some(spaces::BinaryPayload {
            data: vec![env_index as u8, step as u8],
        }),
        route: ModelRouteContext {
            session_id: "test-session".to_string(),
            route_id: "test-route".to_string(),
            request_id: format!("request-{episode_id}-{step}"),
            slots: vec![ModelRouteSlot {
                episode_id: episode_id.to_string(),
                env_index,
                step,
                reset: is_reset,
            }],
        },
        reset: is_reset,
        num_envs: 1,
        env_contract: None,
    }
}

struct SmokeEnv {
    obs_space: spaces::SpaceSpec,
    action_space: spaces::SpaceSpec,
    env_contract: spaces::EnvContract,
}

impl SmokeEnv {
    fn new() -> Self {
        let obs_space = spaces::spaces::BoxSpaceBuilder::scalar(0.0, 255.0, vec![1])
            .dtype(spaces::DType::Uint8)
            .build()
            .unwrap();
        let action_space = spaces::spaces::BoxSpaceBuilder::scalar(0.0, 1.0, vec![1])
            .dtype(spaces::DType::Uint8)
            .build()
            .unwrap();
        let env_contract = spaces::EnvContract {
            id: "SmokeEnv-v0".to_string(),
            observation_space: Some(obs_space.clone()),
            action_space: Some(action_space.clone()),
            metadata: None,
            render_mode: String::new(),
            num_envs: 1,
        };
        Self {
            obs_space,
            action_space,
            env_contract,
        }
    }
}

#[async_trait]
impl crate::SingleEnv for SmokeEnv {
    fn observation_space(&self) -> &spaces::SpaceSpec {
        &self.obs_space
    }

    fn action_space(&self) -> &spaces::SpaceSpec {
        &self.action_space
    }

    fn env_contract(&self) -> &spaces::EnvContract {
        &self.env_contract
    }

    async fn reset(
        &mut self,
        _req: spaces::ResetRequest,
    ) -> std::result::Result<spaces::ResetResult, spaces::EnvRuntimeError> {
        Ok(spaces::ResetResult {
            observation: Some(spaces::SpaceValue::Box(spaces::BoxValue::new(
                vec![0],
                vec![1],
                spaces::DType::Uint8,
            ))),
            info: None,
            episode_id: Some("ep-smoke".to_string()),
        })
    }

    async fn step(
        &mut self,
        _req: spaces::StepRequest,
    ) -> std::result::Result<spaces::StepResult, spaces::EnvRuntimeError> {
        Ok(spaces::StepResult {
            observation: Some(spaces::SpaceValue::Box(spaces::BoxValue::new(
                vec![1],
                vec![1],
                spaces::DType::Uint8,
            ))),
            reward: 1.0,
            terminated: true,
            truncated: false,
            info: None,
        })
    }

    async fn render(
        &mut self,
        _req: spaces::RenderRequest,
    ) -> std::result::Result<spaces::RenderResult, spaces::EnvRuntimeError> {
        Ok(spaces::RenderResult::default())
    }

    async fn close(
        &mut self,
        _req: spaces::CloseRequest,
    ) -> std::result::Result<spaces::CloseResult, spaces::EnvRuntimeError> {
        Ok(spaces::CloseResult)
    }
}

struct SmokeModel {
    predicts: Arc<AtomicUsize>,
    closes: Arc<AtomicUsize>,
}

#[async_trait]
impl ModelHandler for SmokeModel {
    async fn predict(&mut self, _observation: ModelObservation) -> Result<spaces::BinaryPayload> {
        self.predicts.fetch_add(1, Ordering::SeqCst);
        Ok(spaces::BinaryPayload { data: vec![0] })
    }

    async fn on_close(&mut self) -> Result<()> {
        self.closes.fetch_add(1, Ordering::SeqCst);
        Ok(())
    }
}

#[tokio::test]
#[ignore = "requires local socket bind support"]
async fn run_local_smoke_uses_in_process_model() {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let env_address = format!("tcp://{}", listener.local_addr().unwrap());
    let env_server = tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(rlmesh_grpc::env::env_service(
                crate::env::WireEnvAdapter::new(crate::SingleEnvAdapter::new(SmokeEnv::new())),
            ))
            .serve_with_incoming(TcpListenerStream::new(listener))
            .await
            .unwrap()
    });

    let predicts = Arc::new(AtomicUsize::new(0));
    let closes = Arc::new(AtomicUsize::new(0));
    ModelWorker::new(SmokeModel {
        predicts: Arc::clone(&predicts),
        closes: Arc::clone(&closes),
    })
    .run_local_to_async_for_episodes(ConnectAddress::Tcp(env_address), "local-token", 1)
    .await
    .unwrap();

    assert!(predicts.load(Ordering::SeqCst) >= 1);
    assert_eq!(closes.load(Ordering::SeqCst), 1);
    env_server.abort();
}

#[tokio::test]
async fn served_model_configure_route_requires_env_contract() {
    let response = handle_model_request(
        JoinRequest {
            kind: Some(join_request::Kind::ConfigureRoute(ConfigureRouteRequest {
                context: Some(rlmesh_proto::model::v1::PredictContext {
                    session_id: "session-1".to_string(),
                    route_id: "route-1".to_string(),
                    request_id: "configure-1".to_string(),
                    ..Default::default()
                }),
                env_contract: None,
            })),
            request_id: "configure-1".to_string(),
        },
        Arc::new(Mutex::new(SmokeModel {
            predicts: Arc::new(AtomicUsize::new(0)),
            closes: Arc::new(AtomicUsize::new(0)),
        })),
        Arc::new(Mutex::new(HashMap::new())),
        Arc::new(Mutex::new(HashMap::new())),
        None,
    )
    .await;

    assert!(matches!(response.kind, Some(join_response::Kind::Error(_))));
}

#[tokio::test]
async fn served_model_predict_mirrors_route_context() {
    let context = rlmesh_proto::model::v1::PredictContext {
        session_id: "session-1".to_string(),
        route_id: "route-1".to_string(),
        request_id: "request-1".to_string(),
        slots: vec![rlmesh_proto::model::v1::PredictSlot {
            episode_id: "episode-1".to_string(),
            env_index: 0,
            step: 7,
            reset: false,
        }],
    };
    let response = handle_model_request(
        JoinRequest {
            kind: Some(join_request::Kind::Predict(PredictRequest {
                context: Some(context.clone()),
                observation: None,
            })),
            request_id: "request-1".to_string(),
        },
        Arc::new(Mutex::new(SmokeModel {
            predicts: Arc::new(AtomicUsize::new(0)),
            closes: Arc::new(AtomicUsize::new(0)),
        })),
        Arc::new(Mutex::new(HashMap::new())),
        Arc::new(Mutex::new(HashMap::from([(
            "session-1:route-1".to_string(),
            ModelRouteConfig {
                env_contract: None,
                num_envs: 1,
            },
        )]))),
        None,
    )
    .await;

    match response.kind {
        Some(join_response::Kind::Predict(response)) => {
            assert_eq!(response.context, Some(context));
        }
        other => panic!("expected predict response, got {other:?}"),
    }
}

#[tokio::test]
async fn served_model_predict_rejects_route_wider_than_opened_route() {
    let response = handle_model_request(
        JoinRequest {
            kind: Some(join_request::Kind::Predict(PredictRequest {
                context: Some(rlmesh_proto::model::v1::PredictContext {
                    session_id: "session-1".to_string(),
                    route_id: "route-1".to_string(),
                    request_id: "request-1".to_string(),
                    slots: vec![
                        rlmesh_proto::model::v1::PredictSlot {
                            episode_id: "episode-0".to_string(),
                            env_index: 0,
                            ..Default::default()
                        },
                        rlmesh_proto::model::v1::PredictSlot {
                            episode_id: "episode-1".to_string(),
                            env_index: 1,
                            ..Default::default()
                        },
                    ],
                }),
                observation: None,
            })),
            request_id: "request-1".to_string(),
        },
        Arc::new(Mutex::new(SmokeModel {
            predicts: Arc::new(AtomicUsize::new(0)),
            closes: Arc::new(AtomicUsize::new(0)),
        })),
        Arc::new(Mutex::new(HashMap::new())),
        Arc::new(Mutex::new(HashMap::from([(
            "session-1:route-1".to_string(),
            ModelRouteConfig {
                env_contract: None,
                num_envs: 1,
            },
        )]))),
        None,
    )
    .await;

    assert!(matches!(response.kind, Some(join_response::Kind::Error(_))));
}

#[tokio::test]
async fn served_model_close_detaches_and_shutdown_runs_close_hook_once() {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    drop(listener);

    let predicts = Arc::new(AtomicUsize::new(0));
    let closes = Arc::new(AtomicUsize::new(0));
    let server_predicts = Arc::clone(&predicts);
    let server_closes = Arc::clone(&closes);
    let server = tokio::spawn(async move {
        ModelWorker::new(SmokeModel {
            predicts: server_predicts,
            closes: server_closes,
        })
        .serve_to_async_with_options(
            BindAddress::Tcp {
                host: "127.0.0.1".to_string(),
                port,
            },
            "route-token",
            ServeOptions {
                allow_remote_shutdown: true,
                ..ServeOptions::default()
            },
        )
        .await
    });

    let address = format!("tcp://127.0.0.1:{port}");
    let mut client = loop {
        match rlmesh_grpc::ModelClient::connect(&address, "route-token").await {
            Ok(mut client) => {
                client.handshake().await.unwrap();
                break client;
            }
            Err(err) if !server.is_finished() => {
                let _ = err;
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
            Err(err) => panic!("model server did not start: {err}"),
        }
    };

    client.close("client session complete").await.unwrap();
    tokio::time::sleep(Duration::from_millis(50)).await;
    assert!(!server.is_finished());
    assert_eq!(closes.load(Ordering::SeqCst), 0);

    let mut second = rlmesh_grpc::ModelClient::connect(&address, "route-token")
        .await
        .unwrap();
    second.handshake().await.unwrap();
    let shutdown = second.shutdown("test complete").await.unwrap();
    assert!(shutdown.accepted);

    tokio::time::timeout(Duration::from_secs(2), server)
        .await
        .unwrap()
        .unwrap()
        .unwrap();

    assert_eq!(predicts.load(Ordering::SeqCst), 0);
    assert_eq!(closes.load(Ordering::SeqCst), 1);
}

#[tokio::test]
async fn served_lifecycle_allows_predict_for_every_observation() {
    let observations = vec![
        raw_observation("ep-1", 0, 0, true),
        raw_observation("ep-1", 1, 0, false),
    ];
    let mut handler = RecordingHandler::default();
    let mut active_episodes = HashMap::new();

    for observation in observations {
        update_lifecycle(&mut handler, &mut active_episodes, &observation)
            .await
            .unwrap();
        handler.predict(observation).await.unwrap();
    }

    assert_eq!(
        handler.predictions,
        vec![("ep-1".to_string(), 0, 0), ("ep-1".to_string(), 1, 0)]
    );
    assert_eq!(handler.resets, vec![("ep-1".to_string(), 0)]);
}

#[tokio::test]
async fn served_lifecycle_closes_previous_episode_on_reset_boundary() {
    let observations = vec![
        raw_observation("ep-1", 0, 0, true),
        raw_observation("ep-1", 1, 0, false),
        raw_observation("ep-2", 0, 0, true),
        raw_observation("ep-2", 1, 0, false),
    ];
    let mut handler = RecordingHandler::default();
    let mut active_episodes = HashMap::new();

    for observation in observations {
        update_lifecycle(&mut handler, &mut active_episodes, &observation)
            .await
            .unwrap();
    }

    assert_eq!(
        handler.resets,
        vec![("ep-1".to_string(), 0), ("ep-2".to_string(), 0)]
    );
    assert_eq!(
        handler.episode_ends,
        vec![ModelEpisodeEnd {
            episode_id: "ep-1".to_string(),
            env_index: 0,
        }]
    );
}

#[tokio::test]
async fn served_lifecycle_tracks_episode_boundaries_per_env_index() {
    let observations = vec![
        raw_observation("ep-a", 0, 0, true),
        raw_observation("ep-b", 0, 1, true),
        raw_observation("ep-c", 0, 0, true),
    ];
    let mut handler = RecordingHandler::default();
    let mut active_episodes = HashMap::new();

    for observation in observations {
        update_lifecycle(&mut handler, &mut active_episodes, &observation)
            .await
            .unwrap();
    }

    assert_eq!(
        handler.episode_ends,
        vec![ModelEpisodeEnd {
            episode_id: "ep-a".to_string(),
            env_index: 0,
        }]
    );
}

#[tokio::test]
async fn idle_shutdown_waits_for_first_request_then_idle_window() {
    let (tx, mut rx) = mpsc::unbounded_channel();
    let shutdown = tokio::spawn(async move {
        rlmesh_grpc::lifecycle::wait_for_idle_shutdown(&mut rx, Duration::from_millis(25)).await;
    });

    tokio::time::sleep(Duration::from_millis(10)).await;
    assert!(!shutdown.is_finished());

    tx.send(()).unwrap();
    tokio::time::sleep(Duration::from_millis(10)).await;
    tx.send(()).unwrap();
    tokio::time::sleep(Duration::from_millis(10)).await;
    assert!(!shutdown.is_finished());

    tokio::time::timeout(Duration::from_millis(100), shutdown)
        .await
        .unwrap()
        .unwrap();
}
