use axum::http::{HeaderMap, HeaderValue, Method};
use base64::Engine as _;
use mad_hatter::proxy_pass;
use mad_hatter::proxy_pass::CompiledPatterns;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::task::JoinHandle;

// --- helpers ---

/// Start a minimal HTTP echo server that returns method, path, and headers as JSON.
async fn start_echo_server() -> (u16, JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let handle = tokio::spawn(async move {
        loop {
            let accept = listener.accept().await;
            let (mut stream, _) = match accept {
                Ok(v) => v,
                Err(_) => break,
            };
            tokio::spawn(async move {
                let mut buf = vec![0u8; 8192];
                let n = match stream.read(&mut buf).await {
                    Ok(n) if n > 0 => n,
                    _ => return,
                };
                let raw = String::from_utf8_lossy(&buf[..n]);

                // Parse first line
                let first_line = raw.lines().next().unwrap_or("");
                let parts: Vec<&str> = first_line.split_whitespace().collect();
                let method = parts.first().copied().unwrap_or("");
                let path = parts.get(1).copied().unwrap_or("");

                // Collect headers (skip first line, until empty line)
                let mut headers = Vec::new();
                for line in raw.lines().skip(1) {
                    if line.is_empty() {
                        break;
                    }
                    headers.push(line.to_string());
                }

                // Extract body (after \r\n\r\n)
                let body_str = if let Some(pos) = raw.find("\r\n\r\n") {
                    &raw[pos + 4..]
                } else {
                    ""
                };

                // Build JSON response
                let headers_json: Vec<String> = headers
                    .iter()
                    .map(|h| format!("\"{}\"", h.replace('\\', "\\\\").replace('"', "\\\"")))
                    .collect();
                let resp_body = format!(
                    r#"{{"method":"{}","path":"{}","headers":[{}],"body":"{}"}}"#,
                    method,
                    path,
                    headers_json.join(","),
                    body_str.replace('\\', "\\\\").replace('"', "\\\""),
                );

                let response = format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
                    resp_body.len(),
                    resp_body,
                );
                let _ = stream.write_all(response.as_bytes()).await;
            });
        }
    });
    (port, handle)
}

// --- proxy declarations (existing modes) ---

proxy_pass!(SimpleProxy);
proxy_pass!(CookieProxy, cookie_name: "session_token");
proxy_pass!(PrefixProxy, prefix: "/proxy/");
proxy_pass!(PrefixCookieProxy, prefix: "/tunnel/", cookie_name: "my_session");

// --- pattern declarations (new mode) ---

// Mock state for testing
struct MockState {
    proxy_enabled: bool,
}

impl mad_hatter::ProxyResolve for MockState {
    async fn should_proxy(&self, _param: &str) -> bool {
        self.proxy_enabled
    }

    async fn proxy_forward(
        &self,
        _param: &str,
        _request: mad_hatter::HttpTunnelRequest,
    ) -> Option<mad_hatter::HttpTunnelResponse> {
        // Return a simple response
        let mut headers = std::collections::HashMap::new();
        headers.insert("content-type".to_string(), "text/plain".to_string());
        Some(mad_hatter::HttpTunnelResponse {
            status: 200,
            headers,
            body: Some(base64::engine::general_purpose::STANDARD.encode(b"proxied")),
        })
    }
}

proxy_pass! {
    PatternProxy, std::sync::Arc<MockState> {
        patterns: [
            "/api/instances/{instance_id}/",
            "/serve/{instance_id}/",
            "/public/{instance_id}/",
        ],
    }
}

proxy_pass! {
    PatternCookieProxy, std::sync::Arc<MockState> {
        patterns: [
            "/tunnel/{node_id}/",
        ],
        cookie_name: "session_token",
    }
}

// ── Existing tests (direct proxy mode) ──────────────────────

#[tokio::test]
async fn test_proxy_pass_basic_get() {
    let (port, _handle) = start_echo_server().await;
    let proxy = SimpleProxy::new();

    let headers = HeaderMap::new();
    let rest_path = format!("{}/hello/world", port);
    let resp = proxy.forward(&Method::GET, &rest_path, None, &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    assert_eq!(json["method"], "GET");
    assert_eq!(json["path"], "/hello/world");
}

#[tokio::test]
async fn test_proxy_pass_post_with_body() {
    let (port, _handle) = start_echo_server().await;
    let proxy = SimpleProxy::new();

    let mut headers = HeaderMap::new();
    headers.insert("content-type", HeaderValue::from_static("application/json"));
    let rest_path = format!("{}/api/data", port);
    let resp = proxy.forward(&Method::POST, &rest_path, None, &headers, b"hello body").await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    assert_eq!(json["method"], "POST");
    assert_eq!(json["path"], "/api/data");
    assert_eq!(json["body"], "hello body");
}

#[tokio::test]
async fn test_proxy_pass_host_rewrite() {
    let (port, _handle) = start_echo_server().await;
    let proxy = SimpleProxy::new();

    let mut headers = HeaderMap::new();
    headers.insert("host", HeaderValue::from_static("original-host.example.com"));
    let rest_path = format!("{}/check", port);
    let resp = proxy.forward(&Method::GET, &rest_path, None, &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    let hdrs = json["headers"].as_array().unwrap();
    let host_header = hdrs
        .iter()
        .find(|h| h.as_str().unwrap_or("").to_lowercase().starts_with("host:"))
        .unwrap()
        .as_str()
        .unwrap();
    assert!(
        host_header.contains(&format!("localhost:{}", port)),
        "host should be rewritten, got: {}",
        host_header
    );
}

#[tokio::test]
async fn test_proxy_pass_cookie_injection() {
    let (port, _handle) = start_echo_server().await;
    let proxy = CookieProxy::new();

    let mut headers = HeaderMap::new();
    headers.insert(
        axum::http::header::COOKIE,
        HeaderValue::from_static("session_token=my_token_value"),
    );
    let rest_path = format!("{}/auth", port);
    let resp = proxy.forward(&Method::GET, &rest_path, None, &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    let hdrs = json["headers"].as_array().unwrap();
    let cookie_header = hdrs
        .iter()
        .find(|h| h.as_str().unwrap_or("").to_lowercase().starts_with("cookie:"))
        .unwrap()
        .as_str()
        .unwrap();
    assert!(
        cookie_header.contains("session_token=my_token_value"),
        "cookie should contain session_token=my_token_value, got: {}",
        cookie_header
    );
}

#[tokio::test]
async fn test_proxy_pass_no_cookie_when_not_in_headers() {
    let (port, _handle) = start_echo_server().await;
    let proxy = CookieProxy::new();

    let headers = HeaderMap::new();
    let rest_path = format!("{}/no-cookie", port);
    let resp = proxy.forward(&Method::GET, &rest_path, None, &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    let hdrs = json["headers"].as_array().unwrap();
    let has_cookie = hdrs
        .iter()
        .any(|h| h.as_str().unwrap_or("").to_lowercase().starts_with("cookie:"));
    assert!(!has_cookie, "no cookie header should be present");
}

#[tokio::test]
async fn test_proxy_pass_header_passthrough() {
    let (port, _handle) = start_echo_server().await;
    let proxy = SimpleProxy::new();

    let mut headers = HeaderMap::new();
    headers.insert("x-custom-header", HeaderValue::from_static("custom-value"));
    headers.insert("authorization", HeaderValue::from_static("Bearer abc123"));
    let rest_path = format!("{}/headers", port);
    let resp = proxy.forward(&Method::GET, &rest_path, None, &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    let hdrs = json["headers"].as_array().unwrap();

    let has_custom = hdrs.iter().any(|h| h.as_str().unwrap_or("").contains("custom-value"));
    assert!(has_custom, "custom header should be passed through");

    let has_auth = hdrs.iter().any(|h| h.as_str().unwrap_or("").contains("Bearer abc123"));
    assert!(has_auth, "authorization header should be passed through");
}

#[tokio::test]
async fn test_proxy_pass_invalid_port() {
    let proxy = SimpleProxy::new();
    let headers = HeaderMap::new();

    let resp = proxy.forward(&Method::GET, "80/test", None, &headers, &[]).await;
    assert_eq!(resp.status(), axum::http::StatusCode::BAD_REQUEST);
}

#[tokio::test]
async fn test_proxy_pass_502_on_connection_failure() {
    let proxy = SimpleProxy::new();
    let headers = HeaderMap::new();

    let resp = proxy.forward(&Method::GET, "19999/fail", None, &headers, &[]).await;
    assert_eq!(resp.status(), axum::http::StatusCode::BAD_GATEWAY);
}

#[tokio::test]
async fn test_proxy_pass_multiple_methods() {
    let (port, _handle) = start_echo_server().await;
    let proxy = SimpleProxy::new();
    let headers = HeaderMap::new();

    for method in &[Method::GET, Method::POST, Method::PUT, Method::DELETE, Method::PATCH, Method::HEAD, Method::OPTIONS] {
        let rest_path = format!("{}/method-test", port);
        let resp = proxy.forward(method, &rest_path, None, &headers, &[]).await;
        assert_eq!(resp.status(), axum::http::StatusCode::OK, "method {:?} should succeed", method);
    }
}

#[test]
fn test_proxy_pass_prefix_constant() {
    assert_eq!(PrefixProxy::PREFIX, "/proxy/");
    assert_eq!(PrefixCookieProxy::PREFIX, "/tunnel/");
}

#[test]
fn test_proxy_pass_no_prefix_no_constant() {
    let _ = PrefixProxy::PREFIX;
}

#[tokio::test]
async fn test_proxy_pass_query_string() {
    let (port, _handle) = start_echo_server().await;
    let proxy = SimpleProxy::new();
    let headers = HeaderMap::new();

    let rest_path = format!("{}/search", port);
    let resp = proxy.forward(&Method::GET, &rest_path, Some("q=hello&page=1"), &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    assert!(json["path"].as_str().unwrap().contains("q=hello"), "query should be forwarded");
}

// ── Pattern matching tests (new) ────────────────────────────

#[test]
fn test_compiled_patterns_simple() {
    let patterns = CompiledPatterns::compile(&["/api/instances/{instance_id}/"]);
    assert_eq!(
        patterns.extract_first("/api/instances/abc123/"),
        Some("abc123".to_string())
    );
    assert_eq!(
        patterns.extract_first("/api/instances/abc123/foo/bar"),
        Some("abc123".to_string())
    );
    assert_eq!(patterns.extract_first("/other/path"), None);
    assert_eq!(patterns.extract_first("/api/instances/"), None);
}

#[test]
fn test_compiled_patterns_multiple() {
    let patterns = CompiledPatterns::compile(&[
        "/api/instances/{instance_id}/",
        "/serve/{instance_id}/",
        "/public/{instance_id}/",
    ]);
    assert_eq!(
        patterns.extract_first("/api/instances/id1/path"),
        Some("id1".to_string())
    );
    assert_eq!(
        patterns.extract_first("/serve/id2/file.js"),
        Some("id2".to_string())
    );
    assert_eq!(
        patterns.extract_first("/public/id3/apps/test"),
        Some("id3".to_string())
    );
    assert_eq!(patterns.extract_first("/other/path"), None);
    assert_eq!(patterns.extract_first("/api/other/id1/"), None);
}

#[test]
fn test_compiled_patterns_special_chars() {
    // Dots in path should be matched literally
    let patterns = CompiledPatterns::compile(&["/v1.0/items/{id}/"]);
    assert_eq!(
        patterns.extract_first("/v1.0/items/abc/details"),
        Some("abc".to_string())
    );
    // Dot should not match arbitrary char (regex-escaped)
    assert_eq!(patterns.extract_first("/v1X0/items/abc/"), None);
}

#[test]
fn test_compiled_patterns_no_trailing_slash() {
    let patterns = CompiledPatterns::compile(&["/items/{id}"]);
    assert_eq!(
        patterns.extract_first("/items/abc"),
        Some("abc".to_string())
    );
    assert_eq!(
        patterns.extract_first("/items/abc/sub"),
        Some("abc".to_string())
    );
}

#[test]
fn test_macro_extract_param() {
    assert_eq!(
        PatternProxy::extract_param("/api/instances/test123/path"),
        Some("test123".to_string())
    );
    assert_eq!(
        PatternProxy::extract_param("/serve/abc/file"),
        Some("abc".to_string())
    );
    assert_eq!(
        PatternProxy::extract_param("/public/xyz/apps/test"),
        Some("xyz".to_string())
    );
    assert_eq!(PatternProxy::extract_param("/unknown/path"), None);
}

#[test]
fn test_macro_extract_param_with_cookie() {
    assert_eq!(
        PatternCookieProxy::extract_param("/tunnel/node1/rest"),
        Some("node1".to_string())
    );
    assert_eq!(PatternCookieProxy::extract_param("/other/path"), None);
}

// ── Pattern proxy forward test (slave-side) ─────────────────

#[tokio::test]
async fn test_pattern_proxy_forward() {
    let (port, _handle) = start_echo_server().await;
    let proxy = PatternProxy::new();

    let rest_path = format!("{}/hello", port);
    let headers = HeaderMap::new();
    let resp = proxy.forward(&Method::GET, &rest_path, None, &headers, &[]).await;

    let body = axum::body::to_bytes(resp.into_body(), 8192).await.unwrap();
    let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
    assert_eq!(json["method"], "GET");
    assert_eq!(json["path"], "/hello");
}

// ── Middleware integration test ──────────────────────────────

#[tokio::test]
async fn test_proxy_middleware_match_and_forward() {
    use axum::{Router, routing::get};

    let state = std::sync::Arc::new(MockState { proxy_enabled: true });

    // Fallback handler (should NOT be reached for matching paths)
    let app = Router::new()
        .route("/fallback", get(|| async { "fallback" }))
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            PatternProxy::proxy_middleware,
        ))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let client = reqwest::Client::new();

    // Matching path → should be proxied (returns "proxied" from MockState)
    let resp = client
        .get(format!("http://{}/api/instances/abc123/foo", addr))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);
    let body = resp.text().await.unwrap();
    assert_eq!(body, "proxied");
}

#[tokio::test]
async fn test_proxy_middleware_no_match_passthrough() {
    use axum::{Router, routing::get};

    let state = std::sync::Arc::new(MockState { proxy_enabled: true });

    let app = Router::new()
        .route("/fallback", get(|| async { "fallback" }))
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            PatternProxy::proxy_middleware,
        ))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let client = reqwest::Client::new();

    // Non-matching path → should pass through to fallback
    let resp = client
        .get(format!("http://{}/fallback", addr))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);
    let body = resp.text().await.unwrap();
    assert_eq!(body, "fallback");
}

#[tokio::test]
async fn test_proxy_middleware_not_proxied_passthrough() {
    use axum::{Router, routing::get};

    let state = std::sync::Arc::new(MockState { proxy_enabled: false });

    let app = Router::new()
        .route(
            "/api/instances/{instance_id}/data",
            get(|| async { "local" }),
        )
        .layer(axum::middleware::from_fn_with_state(
            state.clone(),
            PatternProxy::proxy_middleware,
        ))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let client = reqwest::Client::new();

    // Matching path but should_proxy returns false → pass through to local handler
    let resp = client
        .get(format!("http://{}/api/instances/abc123/data", addr))
        .send()
        .await
        .unwrap();
    assert_eq!(resp.status(), 200);
    let body = resp.text().await.unwrap();
    assert_eq!(body, "local");
}