mod common;

use common::*;
use rustproxy::RustProxy;
use rustproxy_config::RustProxyOptions;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

#[tokio::test]
async fn test_tcp_forward_echo() {
    let backend_port = next_port();
    let proxy_port = next_port();

    // Start echo backend
    let _backend = start_echo_server(backend_port).await;

    // Configure proxy
    let options = RustProxyOptions {
        routes: vec![make_test_route(proxy_port, None, "127.0.0.1", backend_port)],
        ..Default::default()
    };

    let mut proxy = RustProxy::new(options).unwrap();
    proxy.start().await.unwrap();

    // Wait for proxy to be ready
    assert!(wait_for_port(proxy_port, 2000).await, "Proxy port not ready");

    // Connect and send data
    let result = with_timeout(async {
        let mut stream = TcpStream::connect(format!("127.0.0.1:{}", proxy_port))
            .await
            .unwrap();
        stream.write_all(b"hello world").await.unwrap();

        let mut buf = vec![0u8; 1024];
        let n = stream.read(&mut buf).await.unwrap();
        String::from_utf8_lossy(&buf[..n]).to_string()
    }, 5)
    .await
    .unwrap();

    assert_eq!(result, "hello world");

    proxy.stop().await.unwrap();
}

#[tokio::test]
async fn test_tcp_forward_large_payload() {
    let backend_port = next_port();
    let proxy_port = next_port();

    let _backend = start_echo_server(backend_port).await;

    let options = RustProxyOptions {
        routes: vec![make_test_route(proxy_port, None, "127.0.0.1", backend_port)],
        ..Default::default()
    };

    let mut proxy = RustProxy::new(options).unwrap();
    proxy.start().await.unwrap();
    assert!(wait_for_port(proxy_port, 2000).await);

    let result = with_timeout(async {
        let mut stream = TcpStream::connect(format!("127.0.0.1:{}", proxy_port))
            .await
            .unwrap();

        // Send 1MB of data
        let data = vec![b'A'; 1_000_000];
        stream.write_all(&data).await.unwrap();
        stream.shutdown().await.unwrap();

        // Read all back
        let mut received = Vec::new();
        stream.read_to_end(&mut received).await.unwrap();
        received.len()
    }, 10)
    .await
    .unwrap();

    assert_eq!(result, 1_000_000);

    proxy.stop().await.unwrap();
}

#[tokio::test]
async fn test_tcp_forward_multiple_connections() {
    let backend_port = next_port();
    let proxy_port = next_port();

    let _backend = start_echo_server(backend_port).await;

    let options = RustProxyOptions {
        routes: vec![make_test_route(proxy_port, None, "127.0.0.1", backend_port)],
        ..Default::default()
    };

    let mut proxy = RustProxy::new(options).unwrap();
    proxy.start().await.unwrap();
    assert!(wait_for_port(proxy_port, 2000).await);

    let result = with_timeout(async {
        let mut handles = Vec::new();
        for i in 0..10 {
            let port = proxy_port;
            handles.push(tokio::spawn(async move {
                let mut stream = TcpStream::connect(format!("127.0.0.1:{}", port))
                    .await
                    .unwrap();
                let msg = format!("connection-{}", i);
                stream.write_all(msg.as_bytes()).await.unwrap();

                let mut buf = vec![0u8; 1024];
                let n = stream.read(&mut buf).await.unwrap();
                String::from_utf8_lossy(&buf[..n]).to_string()
            }));
        }

        let mut results = Vec::new();
        for handle in handles {
            results.push(handle.await.unwrap());
        }
        results
    }, 10)
    .await
    .unwrap();

    assert_eq!(result.len(), 10);
    for (i, r) in result.iter().enumerate() {
        assert_eq!(r, &format!("connection-{}", i));
    }

    proxy.stop().await.unwrap();
}

#[tokio::test]
async fn test_tcp_forward_backend_unreachable() {
    let proxy_port = next_port();
    let dead_port = next_port(); // No server on this port

    let options = RustProxyOptions {
        routes: vec![make_test_route(proxy_port, None, "127.0.0.1", dead_port)],
        ..Default::default()
    };

    let mut proxy = RustProxy::new(options).unwrap();
    proxy.start().await.unwrap();
    assert!(wait_for_port(proxy_port, 2000).await);

    // Connection should complete (proxy accepts it) but data should not flow
    let result = with_timeout(async {
        let stream = TcpStream::connect(format!("127.0.0.1:{}", proxy_port)).await;
        stream.is_ok()
    }, 5)
    .await
    .unwrap();

    assert!(result, "Should be able to connect to proxy even if backend is down");

    proxy.stop().await.unwrap();
}

#[tokio::test]
async fn test_tcp_forward_bidirectional() {
    let backend_port = next_port();
    let proxy_port = next_port();

    // Start a prefix echo server to verify data flows in both directions
    let _backend = start_prefix_echo_server(backend_port, "REPLY:").await;

    let options = RustProxyOptions {
        routes: vec![make_test_route(proxy_port, None, "127.0.0.1", backend_port)],
        ..Default::default()
    };

    let mut proxy = RustProxy::new(options).unwrap();
    proxy.start().await.unwrap();
    assert!(wait_for_port(proxy_port, 2000).await);

    let result = with_timeout(async {
        let mut stream = TcpStream::connect(format!("127.0.0.1:{}", proxy_port))
            .await
            .unwrap();
        stream.write_all(b"test data").await.unwrap();

        let mut buf = vec![0u8; 1024];
        let n = stream.read(&mut buf).await.unwrap();
        String::from_utf8_lossy(&buf[..n]).to_string()
    }, 5)
    .await
    .unwrap();

    assert_eq!(result, "REPLY:test data");

    proxy.stop().await.unwrap();
}