use std::sync::Arc;
use std::time::Duration;

use dashmap::DashMap;
use kube::Client;
use kube::api::Api;
use lazy_static::lazy_static;
use tokio::sync::Mutex;
use tokio::task::JoinHandle;
use tokio::time::timeout;
use tokio_util::sync::CancellationToken;
use tracing::instrument;

use crate::kube::listener::{
    ListenerConfig,
    PortForwarder,
    Protocol,
};
use crate::kube::models::{
    PortForward,
    Target,
};
use crate::kube::shared_client::{
    SHARED_CLIENT_MANAGER,
    ServiceClientKey,
};

pub struct PortForwardProcess {
    pub handle: JoinHandle<anyhow::Result<()>>,
    pub direct_forwarder: Option<Arc<PortForwarder>>,
    pub cancellation_token: CancellationToken,
    pub config_id: String,
    pub ws_client_handle: Option<JoinHandle<()>>,
}

impl PortForwardProcess {
    pub fn new(handle: JoinHandle<anyhow::Result<()>>, config_id: String) -> Self {
        Self {
            handle,
            direct_forwarder: None,
            cancellation_token: CancellationToken::new(),
            config_id,
            ws_client_handle: None,
        }
    }

    pub fn new_with_token(
        handle: JoinHandle<anyhow::Result<()>>, config_id: String,
        cancellation_token: CancellationToken,
    ) -> Self {
        Self {
            handle,
            direct_forwarder: None,
            cancellation_token,
            config_id,
            ws_client_handle: None,
        }
    }

    pub fn with_forwarder(
        handle: JoinHandle<anyhow::Result<()>>, forwarder: Arc<PortForwarder>, config_id: String,
    ) -> Self {
        Self {
            handle,
            direct_forwarder: Some(forwarder),
            cancellation_token: CancellationToken::new(),
            config_id,
            ws_client_handle: None,
        }
    }

    pub fn with_forwarder_and_token(
        handle: JoinHandle<anyhow::Result<()>>, forwarder: Arc<PortForwarder>, config_id: String,
        cancellation_token: CancellationToken,
    ) -> Self {
        Self {
            handle,
            direct_forwarder: Some(forwarder),
            cancellation_token,
            config_id,
            ws_client_handle: None,
        }
    }

    pub fn set_ws_client_handle(&mut self, ws_handle: JoinHandle<()>) {
        self.ws_client_handle = Some(ws_handle);
    }

    /// Cleanup and abort the port forward process.
    /// Uses timeouts to prevent blocking on shutdown operations.
    pub async fn cleanup_and_abort(self) {
        const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);
        const CANCEL_PROPAGATION_DELAY: Duration = Duration::from_millis(100);

        tracing::info!("Cancelling port forward for config: {}", self.config_id);

        // Signal cancellation
        self.cancellation_token.cancel();

        // Brief delay for cancellation to propagate (reduced from 1500ms)
        tokio::time::sleep(CANCEL_PROPAGATION_DELAY).await;

        // Shutdown forwarder with timeout to prevent blocking
        if let Some(forwarder) = &self.direct_forwarder {
            tracing::info!(
                "Cleaning up forwarder resources for config: {}",
                self.config_id
            );
            if timeout(SHUTDOWN_TIMEOUT, forwarder.shutdown())
                .await
                .is_err()
            {
                tracing::warn!(
                    "Forwarder shutdown timed out for config: {}, forcing abort",
                    self.config_id
                );
            }
        }

        // Abort the WebSocket client task if it exists
        if let Some(ws_handle) = self.ws_client_handle {
            tracing::info!(
                "Aborting WebSocket client task for config: {}",
                self.config_id
            );
            ws_handle.abort();
        }

        self.handle.abort();
    }

    pub fn cancel(&self) {
        tracing::info!("Cancelling port forward for config: {}", self.config_id);
        self.cancellation_token.cancel();
    }

    pub fn abort(&self) {
        self.handle.abort();
    }

    pub async fn get_current_active_pod(&self) -> Option<String> {
        if let Some(forwarder) = &self.direct_forwarder {
            forwarder.get_current_active_pod().await
        } else {
            None
        }
    }
}

lazy_static! {
    pub static ref CHILD_PROCESSES: DashMap<String, PortForwardProcess> = DashMap::new();
}

impl PortForward {
    pub async fn new(
        target: Target, local_port: impl Into<Option<u16>>,
        local_address: impl Into<Option<String>>, context_name: Option<String>,
        kubeconfig: Option<String>, config_id: i64, workload_type: String,
    ) -> anyhow::Result<Self> {
        let namespace = target.namespace.name_any();

        let client_key = ServiceClientKey::new(context_name.clone(), kubeconfig.clone());
        let client = SHARED_CLIENT_MANAGER.get_client(client_key).await?;

        let shared_client = Client::clone(&client);
        let pod_api = Api::namespaced(shared_client.clone(), &namespace);
        let svc_api = Api::namespaced(shared_client.clone(), &namespace);

        Ok(Self {
            target,
            local_port: local_port.into(),
            local_address: local_address.into(),
            pod_api,
            svc_api,
            client: shared_client,
            context_name: context_name.clone(),
            kubeconfig: kubeconfig.clone(),
            config_id,
            workload_type,
            connection: Arc::new(Mutex::new(None)),
        })
    }

    pub fn local_port(&self) -> u16 {
        self.local_port.unwrap_or(0)
    }

    pub fn local_address(&self) -> Option<String> {
        self.local_address.clone()
    }

    #[instrument(skip(self), fields(config_id = self.config_id))]
    pub async fn cleanup_resources(&self) -> anyhow::Result<()> {
        if let Some(addr) = &self.local_address
            && crate::network_utils::is_custom_loopback_address(addr)
        {
            let _ = crate::network_utils::remove_loopback_address(addr).await;
        }
        Ok(())
    }

    #[instrument(skip(self, tls_acceptor), fields(config_id = self.config_id))]
    pub async fn port_forward_tcp(
        self, tls_acceptor: Option<tokio_rustls::TlsAcceptor>,
    ) -> anyhow::Result<(u16, PortForwardProcess)> {
        let local_addr = self
            .local_address
            .as_deref()
            .unwrap_or("127.0.0.1")
            .to_owned();

        let namespace = self.target.namespace.name_any();

        let mut direct_forwarder = PortForwarder::new(
            &namespace,
            self.target.clone(),
            self.context_name.clone(),
            self.kubeconfig.clone(),
            self.config_id,
        )
        .await?;

        if let Err(e) = direct_forwarder.initialize(&self.target).await {
            direct_forwarder.shutdown().await;
            return Err(e);
        }
        let direct_forwarder = Arc::new(direct_forwarder);

        let listener_config = ListenerConfig {
            local_address: local_addr,
            local_port: self.local_port(),
            protocol: Protocol::Tcp,
            tls_acceptor,
        };

        let forwarder_clone = direct_forwarder.clone();
        let cancellation_token = CancellationToken::new();
        let (port, handle) = match direct_forwarder
            .start_listener(
                listener_config,
                self.config_id,
                self.workload_type.clone(),
                cancellation_token.clone(),
            )
            .await
        {
            Ok(result) => result,
            Err(e) => {
                forwarder_clone.shutdown().await;
                return Err(e);
            }
        };

        let process = PortForwardProcess::with_forwarder_and_token(
            handle,
            forwarder_clone,
            self.config_id.to_string(),
            cancellation_token,
        );
        Ok((port, process))
    }

    pub async fn port_forward_udp(self) -> anyhow::Result<(u16, PortForwardProcess)> {
        let local_addr = self
            .local_address
            .as_deref()
            .unwrap_or("127.0.0.1")
            .to_owned();

        let namespace = self.target.namespace.name_any();

        let mut direct_forwarder = PortForwarder::new(
            &namespace,
            self.target.clone(),
            self.context_name.clone(),
            self.kubeconfig.clone(),
            self.config_id,
        )
        .await?;

        if let Err(e) = direct_forwarder.initialize(&self.target).await {
            direct_forwarder.shutdown().await;
            return Err(e);
        }
        let direct_forwarder = Arc::new(direct_forwarder);

        let listener_config = ListenerConfig {
            local_address: local_addr,
            local_port: self.local_port(),
            protocol: Protocol::Udp,
            tls_acceptor: None,
        };

        let forwarder_clone = direct_forwarder.clone();
        let cancellation_token = CancellationToken::new();
        let (port, handle) = match direct_forwarder
            .start_listener(
                listener_config,
                self.config_id,
                self.workload_type.clone(),
                cancellation_token.clone(),
            )
            .await
        {
            Ok(result) => result,
            Err(e) => {
                forwarder_clone.shutdown().await;
                return Err(e);
            }
        };

        let process = PortForwardProcess::with_forwarder_and_token(
            handle,
            forwarder_clone,
            self.config_id.to_string(),
            cancellation_token,
        );
        Ok((port, process))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use http::{
        Request,
        Response,
        StatusCode,
    };
    use k8s_openapi::List;
    use k8s_openapi::api::core::v1::{
        Pod,
        PodCondition,
        PodSpec,
        PodStatus,
        Service,
        ServicePort,
        ServiceSpec,
    };
    use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
    use kube::Client;
    use kube::client::Body;
    use tokio::net::TcpStream;
    use tokio::time::Duration;
    use tower_test::mock;
    use tracing::info;
    use tracing_subscriber;

    use super::*;
    use crate::kube::models::{
        NameSpace,
        Port,
        TargetSelector,
    };

    fn mock_pod(
        name: &str, namespace: &str, labels: Option<BTreeMap<String, String>>, ready: bool,
    ) -> Pod {
        let status = if ready {
            Some(PodStatus {
                phase: Some("Running".to_string()),
                conditions: Some(vec![PodCondition {
                    type_: "Ready".to_string(),
                    status: "True".to_string(),
                    last_probe_time: None,
                    last_transition_time: None,
                    message: None,
                    observed_generation: None,
                    reason: None,
                }]),
                ..Default::default()
            })
        } else {
            Some(PodStatus {
                phase: Some("Pending".to_string()),
                ..Default::default()
            })
        };
        Pod {
            metadata: ObjectMeta {
                name: Some(name.to_string()),
                namespace: Some(namespace.to_string()),
                labels,
                ..Default::default()
            },
            spec: Some(PodSpec {
                ..Default::default()
            }),
            status,
        }
    }

    fn mock_service(
        name: &str, namespace: &str, selector: Option<BTreeMap<String, String>>,
    ) -> Service {
        Service {
            metadata: ObjectMeta {
                name: Some(name.to_string()),
                namespace: Some(namespace.to_string()),
                ..Default::default()
            },
            spec: Some(ServiceSpec {
                selector,
                ports: Some(vec![ServicePort {
                    port: 80,
                    ..Default::default()
                }]),
                ..Default::default()
            }),
            ..Default::default()
        }
    }

    #[tokio::test]
    async fn test_child_processes_map() {
        CHILD_PROCESSES.clear();

        assert_eq!(CHILD_PROCESSES.len(), 0);

        CHILD_PROCESSES.insert(
            "test-key".to_string(),
            PortForwardProcess::new(dummy_handle(), "test-key".to_string()),
        );
        assert_eq!(CHILD_PROCESSES.len(), 1);
        assert!(CHILD_PROCESSES.contains_key("test-key"));

        let handle = CHILD_PROCESSES.remove("test-key");
        if let Some((_, h)) = handle {
            h.abort();
        }
        assert_eq!(CHILD_PROCESSES.len(), 0);
    }

    fn dummy_handle() -> JoinHandle<anyhow::Result<()>> {
        use std::pin::Pin;
        use std::task::{
            Context,
            Poll,
        };

        struct DummyFuture;
        impl std::future::Future for DummyFuture {
            type Output = anyhow::Result<()>;
            fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
                Poll::Ready(Ok(()))
            }
        }

        tokio::task::spawn(DummyFuture)
    }

    #[tokio::test]
    async fn test_process_cancellation() {
        let process = PortForwardProcess::new(dummy_handle(), "test-process".to_string());
        let cancellation_token = process.cancellation_token.clone();

        let task = tokio::spawn(async move {
            tokio::select! {
                _ = cancellation_token.cancelled() => true,
                _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => false,
            }
        });

        process.cancel();

        assert!(task.await.unwrap());
    }

    #[tokio::test]
    async fn test_local_port_and_address() {
        let target = Target {
            selector: TargetSelector::ServiceName("test-service".to_string()),
            port: Port::Number(8080),
            namespace: NameSpace(Some("default".to_string())),
        };

        let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
        let client = Client::new(mock_service, "default");

        let pf_with_port = PortForward {
            target: target.clone(),
            local_port: Some(9090),
            local_address: Some("127.0.0.1".to_string()),
            pod_api: Api::namespaced(client.clone(), "default"),
            svc_api: Api::namespaced(client.clone(), "default"),
            client: client.clone(),
            context_name: None,
            kubeconfig: None,
            config_id: 1,
            workload_type: "pod".to_string(),
            connection: Arc::new(Mutex::new(None)),
        };

        assert_eq!(pf_with_port.local_port(), 9090);
        assert_eq!(pf_with_port.local_address(), Some("127.0.0.1".to_string()));

        let pf_defaults = PortForward {
            target: target.clone(),
            local_port: Some(0),
            local_address: None,
            pod_api: Api::namespaced(client.clone(), "default"),
            svc_api: Api::namespaced(client.clone(), "default"),
            client: client.clone(),
            context_name: None,
            kubeconfig: None,
            config_id: 2,
            workload_type: "pod".to_string(),
            connection: Arc::new(Mutex::new(None)),
        };

        assert_eq!(pf_defaults.local_port(), 0);
        assert_eq!(pf_defaults.local_address(), None);
    }

    #[tokio::test]
    async fn test_port_forward_new_context_propagation() {
        let target = Target {
            selector: TargetSelector::ServiceName("test-service".to_string()),
            port: Port::Number(8080),
            namespace: NameSpace(Some("default".to_string())),
        };
        let context_name = Some("my-kube-context".to_string());
        let _kubeconfig = Some("/path/to/config".to_string());

        let (mock_service, _handle) = mock::pair::<Request<Body>, Response<Body>>();
        let client = Client::new(mock_service, "default");

        let pf = PortForward {
            target,
            local_port: Some(12345),
            local_address: None,
            pod_api: Api::namespaced(client.clone(), "default"),
            svc_api: Api::namespaced(client.clone(), "default"),
            client: client.clone(),
            context_name: context_name.clone(),
            kubeconfig: None,
            config_id: 1,
            workload_type: "service".to_string(),
            connection: Arc::new(Mutex::new(None)),
        };

        assert_eq!(pf.context_name, context_name);
    }

    fn setup_mock_pf_for_api_test() -> (
        PortForward,
        Client,
        mock::Handle<Request<Body>, Response<Body>>,
    ) {
        let target = Target {
            selector: TargetSelector::ServiceName("test-svc".to_string()),
            port: Port::Number(80),
            namespace: NameSpace(Some("test-ns".to_string())),
        };

        let (mock_service, handle) = mock::pair::<Request<Body>, Response<Body>>();
        let client = Client::new(mock_service, "test-ns");

        let pf = PortForward {
            target,
            local_port: Some(0),
            local_address: None,
            pod_api: Api::namespaced(client.clone(), "test-ns"),
            svc_api: Api::namespaced(client.clone(), "test-ns"),
            client: client.clone(),
            context_name: None,
            kubeconfig: None,
            config_id: 1,
            workload_type: "service".to_string(),
            connection: Arc::new(Mutex::new(None)),
        };
        (pf, client, handle)
    }

    async fn mock_kube_api_calls(handle: &mut mock::Handle<Request<Body>, Response<Body>>) {
        info!("Mock server: Starting to handle requests");

        let mut service_requests = 0;
        let mut pod_requests = 0;
        let mut portforward_requests = 0;

        for i in 0..10 {
            info!("Mock server: Expecting request {}", i + 1);

            let result =
                tokio::time::timeout(Duration::from_millis(100), handle.next_request()).await;

            let (request, send) = match result {
                Ok(Some((req, send))) => (req, send),
                Ok(None) => {
                    info!("Mock server: No more requests");
                    break;
                }
                Err(_) => {
                    info!(
                        "Mock server: Timeout waiting for request, checking if we have minimum required"
                    );
                    if service_requests > 0 && pod_requests > 0 {
                        info!("Mock server: Have minimum required requests, stopping");
                        break;
                    }
                    continue;
                }
            };

            info!(
                "Mock server: Received request for path: {}",
                request.uri().path()
            );

            if request.uri().path().contains("/services/") {
                service_requests += 1;
                info!("Mock server: Handling GET Service (#{service_requests})");
                assert_eq!(request.method(), "GET");
                let svc = mock_service(
                    "test-svc",
                    "test-ns",
                    Some(
                        [("app".to_string(), "my-app".to_string())]
                            .into_iter()
                            .collect(),
                    ),
                );
                let response = Response::builder()
                    .status(200)
                    .body(Body::from(serde_json::to_vec(&svc).unwrap()))
                    .unwrap();
                info!("Mock server: Sending service response");
                send.send_response(response);
            } else if request.uri().path().contains("/pods")
                && !request.uri().path().contains("/portforward")
            {
                pod_requests += 1;
                info!("Mock server: Handling LIST Pods (#{pod_requests})");
                assert_eq!(request.method(), "GET");

                let pod = mock_pod(
                    "test-pod-123",
                    "test-ns",
                    Some(
                        [("app".to_string(), "my-app".to_string())]
                            .into_iter()
                            .collect(),
                    ),
                    true,
                );
                let pod_list: List<Pod> = List {
                    items: vec![pod],
                    ..Default::default()
                };
                let response = Response::builder()
                    .status(200)
                    .body(Body::from(serde_json::to_vec(&pod_list).unwrap()))
                    .unwrap();
                info!("Mock server: Sending pods response");
                send.send_response(response);
            } else if request.uri().path().contains("/portforward") {
                portforward_requests += 1;
                info!("Mock server: Handling GET Portforward (#{portforward_requests})");
                assert_eq!(request.method(), "GET");

                let response = Response::builder()
                    .status(StatusCode::SWITCHING_PROTOCOLS)
                    .header(http::header::UPGRADE, "websocket")
                    .header(http::header::CONNECTION, "Upgrade")
                    .header(http::header::SEC_WEBSOCKET_ACCEPT, "dummy_accept_key")
                    .body(Body::empty())
                    .unwrap();
                info!("Mock server: Sending portforward response");
                send.send_response(response);

                if portforward_requests >= 1 {
                    info!("Mock server: Got portforward request, can exit");
                    break;
                }
            }
        }

        info!(
            "Mock server: Handled {service_requests} service, {pod_requests} pod, {portforward_requests} portforward requests"
        );
    }

    #[tokio::test]
    async fn test_port_forward_tcp_api_calls() {
        let _ = tracing_subscriber::fmt().with_test_writer().try_init();
        info!("Starting test_port_forward_tcp_api_calls");

        let (pf_base, _client, _handle) = setup_mock_pf_for_api_test();

        let (mock_service_test, mut handle_test) = mock::pair::<Request<Body>, Response<Body>>();
        let client_test = Client::new(mock_service_test, "test-ns");

        let target_clone = pf_base.target.clone();

        let pf_test = PortForward {
            target: target_clone,
            local_port: Some(0),
            local_address: Some("127.0.0.1".to_string()),
            pod_api: Api::namespaced(client_test.clone(), "test-ns"),
            svc_api: Api::namespaced(client_test.clone(), "test-ns"),
            client: client_test.clone(),
            context_name: None,
            kubeconfig: None,
            config_id: 1,
            workload_type: "service".to_string(),
            connection: Arc::new(Mutex::new(None)),
        };

        info!("Spawning mock server task");
        let mock_server_task = tokio::spawn(async move {
            mock_kube_api_calls(&mut handle_test).await;
        });

        info!("Calling port_forward_tcp");
        let pf_result = pf_test.port_forward_tcp(None).await;
        if pf_result.is_err() {
            info!(
                "Port forward failed as expected in test environment: {:?}",
                pf_result.err()
            );
            return;
        }
        let (bound_port, server_task_handle) = pf_result.unwrap();
        assert_ne!(bound_port, 0, "Listener did not bind to a dynamic port");

        info!("Simulating client connection");
        let connect_addr = format!("127.0.0.1:{bound_port}");
        let connect_task = tokio::spawn(async move {
            match TcpStream::connect(&connect_addr).await {
                Ok(stream) => {
                    drop(stream);
                    Ok(())
                }
                Err(e) => Err(e),
            }
        });

        let connect_result = tokio::time::timeout(Duration::from_secs(1), connect_task).await;
        match connect_result {
            Ok(Ok(Ok(()))) => {
                info!("Client connection succeeded");
            }
            Ok(Ok(Err(e))) => {
                info!(
                    "Client connection failed as expected in test environment: {}",
                    e
                );
            }
            Ok(Err(_)) | Err(_) => {
                info!("Client connection timed out or failed as expected in test environment");
            }
        }
        info!("Client connection simulated");

        info!("Waiting for mock server task");
        let mock_server_result =
            tokio::time::timeout(Duration::from_secs(5), mock_server_task).await;
        assert!(mock_server_result.is_ok(), "Mock server task timed out");
        assert!(
            mock_server_result.unwrap().is_ok(),
            "Mock server task failed"
        );
        info!("Mock server task completed");

        server_task_handle.abort();
        info!("Finished test_port_forward_tcp_api_calls");
    }

    #[tokio::test]
    async fn test_port_forward_udp_api_calls() {
        let _ = tracing_subscriber::fmt().with_test_writer().try_init();
        info!("Starting test_port_forward_udp_api_calls");

        let (pf, _client, mut handle) = setup_mock_pf_for_api_test();

        info!("Spawning mock server task");
        let mock_server = tokio::spawn(async move {
            mock_kube_api_calls(&mut handle).await;
        });

        info!("Calling port_forward_udp");
        let result = tokio::time::timeout(Duration::from_secs(5), pf.port_forward_udp()).await;
        info!("port_forward_udp call returned");

        info!("Waiting for mock server task");
        let mock_server_result = mock_server.await;
        assert!(
            mock_server_result.is_ok(),
            "Mock server did not handle all requests"
        );

        assert!(result.is_ok(), "port_forward_udp timed out unexpectedly");
        assert!(
            result.unwrap().is_err(),
            "port_forward_udp succeeded unexpectedly, expected error after API calls"
        );
        info!("Finished test_port_forward_udp_api_calls");
    }
}
