use mad_hatter::tunnel_service;
use mad_hatter::tunnel::{TunnelEndpoint, Dispatch, TunnelCaller};
use mad_hatter::tunnel_client::DefaultCodec;
use mad_hatter::tunnel_bind;
use serde::{Serialize, Deserialize};
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use futures_util::{SinkExt, StreamExt};
use tokio_tungstenite::tungstenite::Message;

// ─── Service Definition ───

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct Info {
    value: String,
}

#[tunnel_service]
trait CalcService: Send + Sync {
    async fn multiply(&self, a: i32, b: i32) -> Result<i32, String>;
    async fn greet(&self, name: String) -> Result<String, String>;
    async fn get_info(&self) -> Result<Info, String>;
}

struct CalcHandler;

#[async_trait]
impl CalcService for CalcHandler {
    async fn multiply(&self, a: i32, b: i32) -> Result<i32, String> {
        Ok(a * b)
    }

    async fn greet(&self, name: String) -> Result<String, String> {
        Ok(format!("hello, {}!", name))
    }

    async fn get_info(&self) -> Result<Info, String> {
        Ok(Info { value: "bind-test".into() })
    }
}

fn make_dispatchers() -> Vec<Box<dyn Dispatch>> {
    vec![CalcServiceDispatcher::boxed(Arc::new(CalcHandler))]
}

// ─── WS Server Helper ───

async fn start_ws_server() -> (std::net::SocketAddr, mpsc::Sender<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let (shutdown_tx, mut shutdown_rx) = mpsc::channel::<()>(1);

    tokio::spawn(async move {
        loop {
            tokio::select! {
                accept_result = listener.accept() => {
                    match accept_result {
                        Ok((stream, _)) => {
                            let ws_stream = match tokio_tungstenite::accept_async(stream).await {
                                Ok(ws) => ws,
                                Err(_) => continue,
                            };
                            serve_ws_connection(ws_stream);
                        }
                        Err(_) => break,
                    }
                }
                _ = shutdown_rx.recv() => {
                    break;
                }
            }
        }
    });

    (addr, shutdown_tx)
}

fn serve_ws_connection<S>(ws_stream: S)
where
    S: futures_util::Stream<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
        + futures_util::Sink<Message, Error = tokio_tungstenite::tungstenite::Error>
        + Unpin
        + Send
        + 'static,
{
    let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();
    let (outgoing_tx, outgoing_rx) = mpsc::unbounded_channel();

    let _endpoint = TunnelEndpoint::new(
        incoming_rx,
        outgoing_tx,
        make_dispatchers(),
        Duration::from_secs(5),
    );

    let (mut ws_sink, mut ws_read) = ws_stream.split();

    // Reader: WS → incoming_tx
    tokio::spawn(async move {
        while let Some(result) = ws_read.next().await {
            match result {
                Ok(Message::Text(text)) => {
                    if incoming_tx.send(text.to_string()).is_err() {
                        break;
                    }
                }
                Ok(Message::Close(_)) | Err(_) => break,
                _ => {}
            }
        }
    });

    // Writer: outgoing_rx → WS
    tokio::spawn(async move {
        let mut outgoing_rx = outgoing_rx;
        while let Some(text) = outgoing_rx.recv().await {
            if ws_sink.send(Message::Text(text.into())).await.is_err() {
                break;
            }
        }
    });
}

// ─── Tunnel Bind Declarations ───

// Minimal form
tunnel_bind!(BasicTunnel, "ws://127.0.0.1:9999/tunnel");

// With custom timeouts
tunnel_bind!(CustomTunnel, "ws://127.0.0.1:9999/tunnel",
    heartbeat_secs: 15,
    call_timeout_secs: 10,
);

// With env_host
tunnel_bind!(EnvTunnel, "ws://127.0.0.1:9999/tunnel",
    env_host: "TUNNEL_BIND_INTEGRATION_HOST",
    heartbeat_secs: 20,
    call_timeout_secs: 15,
);

// With all options in different order
tunnel_bind!(FullTunnel, "ws://127.0.0.1:9999/tunnel",
    call_timeout_secs: 25,
    env_host: "TUNNEL_BIND_FULL_HOST",
    heartbeat_secs: 45,
);

// ─── Tests ───

#[tokio::test]
async fn test_tunnel_bind_basic_rpc() {
    let (addr, _shutdown) = start_ws_server().await;
    std::env::set_var("TUNNEL_BIND_INTEGRATION_HOST", addr.to_string());

    let tunnel = EnvTunnel::connect(DefaultCodec, make_dispatchers(), vec![])
        .await
        .expect("connect should succeed");

    let proxy = CalcServiceProxy::new(tunnel.clone());
    let result = proxy.multiply(6, 7).await.unwrap();
    assert_eq!(result, 42);

    let greeting = proxy.greet("world".to_string()).await.unwrap();
    assert_eq!(greeting, "hello, world!");

    tunnel.shutdown();
    std::env::remove_var("TUNNEL_BIND_INTEGRATION_HOST");
}

#[tokio::test]
async fn test_tunnel_bind_complex_return() {
    let (addr, _shutdown) = start_ws_server().await;
    std::env::set_var("TUNNEL_BIND_FULL_HOST", addr.to_string());

    let tunnel = FullTunnel::connect(DefaultCodec, make_dispatchers(), vec![])
        .await
        .expect("connect should succeed");

    let proxy = CalcServiceProxy::new(tunnel.clone());
    let info = proxy.get_info().await.unwrap();
    assert_eq!(info, Info { value: "bind-test".into() });

    tunnel.shutdown();
    std::env::remove_var("TUNNEL_BIND_FULL_HOST");
}

#[tokio::test]
async fn test_tunnel_bind_as_tunnel_caller() {
    // Verify that the generated struct implements TunnelCaller (compile-time check)
    // and can be used as Arc<dyn TunnelCaller>
    let (addr, _shutdown) = start_ws_server().await;
    std::env::set_var("TUNNEL_BIND_CALLER_HOST", addr.to_string());

    tunnel_bind!(CallerTunnel, "ws://127.0.0.1:1/tunnel",
        env_host: "TUNNEL_BIND_CALLER_HOST",
    );

    let tunnel = CallerTunnel::connect(DefaultCodec, make_dispatchers(), vec![])
        .await
        .expect("connect should succeed");

    // Verify Arc<CallerTunnel> can be used as Arc<dyn TunnelCaller>
    let caller: Arc<dyn TunnelCaller> = tunnel.clone();

    // Use via proxy (which takes Arc<dyn TunnelCaller>)
    let proxy = CalcServiceProxy::new(caller);
    let result = proxy.multiply(6, 7).await.unwrap();
    assert_eq!(result, 42);

    tunnel.shutdown();
    std::env::remove_var("TUNNEL_BIND_CALLER_HOST");
}

#[tokio::test]
async fn test_tunnel_bind_shutdown() {
    let (addr, _shutdown) = start_ws_server().await;
    std::env::set_var("TUNNEL_BIND_INTEGRATION_HOST", addr.to_string());

    let tunnel = EnvTunnel::connect(DefaultCodec, make_dispatchers(), vec![])
        .await
        .expect("connect should succeed");

    // RPC works before shutdown
    let proxy = CalcServiceProxy::new(tunnel.clone());
    assert_eq!(proxy.multiply(2, 3).await.unwrap(), 6);

    // Shutdown
    tunnel.shutdown();
    tokio::time::sleep(Duration::from_millis(100)).await;

    // send_raw should fail after shutdown
    let result = tunnel.send_raw("test".to_string());
    assert!(result.is_err());

    std::env::remove_var("TUNNEL_BIND_INTEGRATION_HOST");
}

#[tokio::test]
async fn test_tunnel_bind_send_raw() {
    let (addr, _shutdown) = start_ws_server().await;
    std::env::set_var("TUNNEL_BIND_INTEGRATION_HOST", addr.to_string());

    let tunnel = EnvTunnel::connect(DefaultCodec, make_dispatchers(), vec![])
        .await
        .expect("connect should succeed");

    // send_raw should succeed (message won't be RPC, server ignores it)
    let result = tunnel.send_raw("raw_protocol_message".to_string());
    assert!(result.is_ok());

    tunnel.shutdown();
    std::env::remove_var("TUNNEL_BIND_INTEGRATION_HOST");
}

#[tokio::test]
async fn test_tunnel_bind_client_accessor() {
    let (addr, _shutdown) = start_ws_server().await;
    std::env::set_var("TUNNEL_BIND_INTEGRATION_HOST", addr.to_string());

    let tunnel = EnvTunnel::connect(DefaultCodec, make_dispatchers(), vec![])
        .await
        .expect("connect should succeed");

    // client() should return Arc<TunnelClient>
    let client = tunnel.client();
    // Can create proxy from client directly
    let proxy = CalcServiceProxy::new(client.clone());
    assert_eq!(proxy.multiply(5, 5).await.unwrap(), 25);

    tunnel.shutdown();
    std::env::remove_var("TUNNEL_BIND_INTEGRATION_HOST");
}

#[tokio::test]
async fn test_tunnel_bind_macro_compiles_all_forms() {
    // Verify that all declared tunnel_bind! forms compile and have expected API.
    // BasicTunnel, CustomTunnel, EnvTunnel, FullTunnel are all declared above.
    fn _assert_caller<T: TunnelCaller>() {}
    _assert_caller::<BasicTunnel>();
    _assert_caller::<CustomTunnel>();
    _assert_caller::<EnvTunnel>();
    _assert_caller::<FullTunnel>();
}