pub mod auth;
pub mod config;

use std::collections::HashMap;
use std::sync::Arc;

use axum::Router;
use axum::body::Bytes;
use axum::extract::DefaultBodyLimit;
use axum::extract::ws::{Message, WebSocket};
use axum::extract::{State, WebSocketUpgrade};
use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::any;
use base64::Engine;
use dashmap::DashMap;
use futures_util::{SinkExt, StreamExt};
use sha2::{Digest, Sha256};
use tokio::sync::{RwLock, oneshot};

use crate::protocol::{
    RegisterAck, RegisterRequest, SubdomainOffer, SubdomainPick, TunnelRequest, TunnelResponse,
    is_hop_by_hop,
};
use auth::{AuthError, AuthProviderConfig, subdomain_matches, verify_token};
use config::RelayConfig;

// ── Relay server ────────────────────────────────────────────────────────

type PendingRequests = Arc<RwLock<HashMap<String, oneshot::Sender<TunnelResponse>>>>;
type TunnelSender = tokio::sync::mpsc::Sender<String>;

struct TunnelConnection {
    sender: TunnelSender,
    pending: PendingRequests,
    /// Signal to notify the tunnel handler it has been evicted.
    evict: tokio::sync::Notify,
}

struct RelayState {
    /// subdomain → active tunnel connection
    tunnels: DashMap<String, Arc<TunnelConnection>>,
    /// Base domain for tunnel URLs
    base_domain: String,
    /// Auth mode for token validation
    auth: AuthMode,
    /// Max tunnel response body size (decoded bytes)
    max_response_body: usize,
}

/// How the relay validates tunnel registration tokens.
enum AuthMode {
    /// No authentication — anyone can tunnel
    Open,
    /// Static tokens from config file (token → allowed subdomain patterns)
    Static(HashMap<String, Vec<String>>),
    /// External auth provider API
    Provider(AuthProviderConfig),
}

/// Derive a consistent subdomain from a token.
fn token_to_subdomain(token: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(token.as_bytes());
    let hash = hasher.finalize();
    hex::encode(&hash[..6])
}

mod hex {
    pub fn encode(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{b:02x}")).collect()
    }
}

/// Build the relay Axum app without binding or serving.
///
/// Returns `(router, port)` — the caller controls the TCP listener and shutdown.
pub fn build_relay_app(cfg: RelayConfig) -> (Router, u16) {
    let auth = if !cfg.tokens.is_empty() {
        let count = cfg.tokens.len();
        println!(
            "  {} static tokens: {} token(s) configured",
            colored::Colorize::green("✓"),
            count,
        );
        AuthMode::Static(cfg.tokens)
    } else if let Some(url) = cfg.auth_provider {
        let secret = cfg
            .auth_provider_secret
            .expect("auth_provider_secret is required when auth_provider is set");
        println!("  {} auth provider enabled", colored::Colorize::green("✓"));
        AuthMode::Provider(AuthProviderConfig {
            url: url.trim_end_matches('/').to_string(),
            secret,
            client: reqwest::Client::builder()
                .timeout(std::time::Duration::from_secs(5))
                .build()
                .unwrap(),
        })
    } else {
        println!(
            "  {} open mode (anyone can tunnel)",
            colored::Colorize::yellow("!"),
        );
        AuthMode::Open
    };

    const DEFAULT_MAX_RESPONSE_BODY_SIZE: usize = 10 * 1024 * 1024;
    let max_response_body = cfg
        .max_response_body_size
        .unwrap_or(DEFAULT_MAX_RESPONSE_BODY_SIZE);

    let state = Arc::new(RelayState {
        tunnels: DashMap::new(),
        base_domain: cfg.relay_domain,
        auth,
        max_response_body,
    });

    const DEFAULT_MAX_REQUEST_BODY_SIZE: usize = 5 * 1024 * 1024;
    let max_body = cfg
        .max_request_body_size
        .unwrap_or(DEFAULT_MAX_REQUEST_BODY_SIZE);

    let app = Router::new()
        .route("/_tunnel/register", any(handle_register))
        .fallback(any(handle_tunnel_request))
        .with_state(state)
        .layer(DefaultBodyLimit::max(max_body));

    let port = cfg.port;
    (app, port)
}

/// Start the relay server (binds and serves until process exits).
pub async fn start_relay(cfg: RelayConfig) {
    let (app, port) = build_relay_app(cfg);

    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{port}"))
        .await
        .expect("Failed to bind relay");

    println!(
        "  {} relay listening on :{port}",
        colored::Colorize::green("mcpr")
    );

    axum::serve(listener, app).await.expect("Relay failed");
}

/// WebSocket registration endpoint.
/// Token is sent as the first message after upgrade (not in query string).
async fn handle_register(ws: WebSocketUpgrade, State(state): State<Arc<RelayState>>) -> Response {
    ws.on_upgrade(move |socket| handle_tunnel_ws(socket, state))
}

async fn handle_tunnel_ws(socket: WebSocket, state: Arc<RelayState>) {
    let (mut ws_sink, mut ws_stream) = socket.split();

    // Read RegisterRequest as the first message (contains token + optional subdomain)
    let reg: RegisterRequest = loop {
        match ws_stream.next().await {
            Some(Ok(Message::Text(text))) => match serde_json::from_str(&text) {
                Ok(req) => break req,
                Err(_) => continue,
            },
            Some(Err(_)) | None => return,
            _ => continue,
        }
    };

    let token = reg.token;
    let requested_subdomain = reg.subdomain;

    // Helper to close with error
    async fn close_with_error(
        ws_sink: &mut futures_util::stream::SplitSink<WebSocket, Message>,
        reason: &str,
    ) {
        let _ = ws_sink
            .send(Message::Close(Some(axum::extract::ws::CloseFrame {
                code: 4001,
                reason: reason.into(),
            })))
            .await;
    }

    // Helper to offer subdomains and wait for client pick
    async fn offer_and_pick(
        ws_sink: &mut futures_util::stream::SplitSink<WebSocket, Message>,
        ws_stream: &mut futures_util::stream::SplitStream<WebSocket>,
        allowed: &[String],
    ) -> Option<String> {
        let offer = SubdomainOffer {
            subdomains: allowed.to_vec(),
        };
        if ws_sink
            .send(Message::Text(serde_json::to_string(&offer).unwrap().into()))
            .await
            .is_err()
        {
            return None;
        }
        // Wait for SubdomainPick
        loop {
            match ws_stream.next().await {
                Some(Ok(Message::Text(text))) => {
                    if let Ok(pick) = serde_json::from_str::<SubdomainPick>(&text) {
                        return Some(pick.subdomain);
                    }
                    continue;
                }
                _ => return None,
            }
        }
    }

    // Resolve subdomain: if requested and allowed → use it; if not allowed or missing → offer pick
    async fn resolve_subdomain(
        requested: Option<String>,
        allowed: &[String],
        ws_sink: &mut futures_util::stream::SplitSink<WebSocket, Message>,
        ws_stream: &mut futures_util::stream::SplitStream<WebSocket>,
    ) -> Option<String> {
        // If requested and it matches, use it directly
        if let Some(ref sub) = requested
            && subdomain_matches(allowed, sub)
        {
            return Some(sub.clone());
        }
        // Auto-assign if exactly one concrete subdomain
        if allowed.len() == 1 && !allowed[0].contains('*') {
            return Some(allowed[0].clone());
        }
        // Otherwise offer the list and let the client pick
        let picked = offer_and_pick(ws_sink, ws_stream, allowed).await?;
        if subdomain_matches(allowed, &picked) {
            Some(picked)
        } else {
            close_with_error(
                ws_sink,
                &format!("subdomain '{}' not authorized for this token", picked),
            )
            .await;
            None
        }
    }

    // Validate token and resolve subdomain based on auth mode
    let subdomain = match &state.auth {
        AuthMode::Open => Some(requested_subdomain.unwrap_or_else(|| token_to_subdomain(&token))),
        AuthMode::Static(tokens) => match tokens.get(&token) {
            Some(allowed) => {
                resolve_subdomain(requested_subdomain, allowed, &mut ws_sink, &mut ws_stream).await
            }
            None => {
                close_with_error(&mut ws_sink, "invalid token").await;
                return;
            }
        },
        AuthMode::Provider(auth) => {
            let sub_for_verify = requested_subdomain.as_deref().unwrap_or("");
            match verify_token(auth, &token, sub_for_verify).await {
                Ok(allowed) => {
                    resolve_subdomain(requested_subdomain, &allowed, &mut ws_sink, &mut ws_stream)
                        .await
                }
                Err(AuthError::InvalidToken(msg)) => {
                    close_with_error(&mut ws_sink, &msg).await;
                    return;
                }
                Err(AuthError::ProviderUnavailable(msg)) => {
                    println!(
                        "  {} auth provider error: {}",
                        colored::Colorize::red("✗"),
                        msg
                    );
                    close_with_error(&mut ws_sink, "auth provider unavailable").await;
                    return;
                }
            }
        }
    };

    let subdomain = match subdomain {
        Some(s) => s,
        None => return,
    };

    let url = format!("https://{}.{}", subdomain, state.base_domain);

    // Send registration ack
    let ack = RegisterAck {
        subdomain: subdomain.clone(),
        url: url.clone(),
    };
    if ws_sink
        .send(Message::Text(serde_json::to_string(&ack).unwrap().into()))
        .await
        .is_err()
    {
        return;
    }

    // Create channel for sending requests to this tunnel client
    let (tx, mut rx) = tokio::sync::mpsc::channel::<String>(64);
    let pending: PendingRequests = Arc::new(RwLock::new(HashMap::new()));

    let conn = Arc::new(TunnelConnection {
        sender: tx,
        pending: pending.clone(),
        evict: tokio::sync::Notify::new(),
    });

    // Evict existing tunnel on same subdomain (if any) before registering
    if let Some((_, old)) = state.tunnels.remove(&subdomain) {
        old.evict.notify_one();
        println!(
            "  {} evicted old tunnel: {}",
            colored::Colorize::yellow("⇄"),
            subdomain
        );
    }

    // Register tunnel
    let conn_for_evict = conn.clone();
    state.tunnels.insert(subdomain.clone(), conn);

    println!(
        "  {} tunnel registered: {}",
        colored::Colorize::green("↑"),
        colored::Colorize::cyan(url.as_str())
    );

    // Spawn task to forward outbound messages (relay → client) from channel to WS.
    // Also listens for eviction signal to send a close frame with reason.
    let send_task = tokio::spawn(async move {
        loop {
            tokio::select! {
                msg = rx.recv() => {
                    match msg {
                        Some(msg) => {
                            if ws_sink.send(Message::Text(msg.into())).await.is_err() {
                                break;
                            }
                        }
                        None => break,
                    }
                }
                () = conn_for_evict.evict.notified() => {
                    let _ = ws_sink
                        .send(Message::Close(Some(axum::extract::ws::CloseFrame {
                            code: 4002,
                            reason: "evicted: another client registered with the same tunnel".into(),
                        })))
                        .await;
                    break;
                }
            }
        }
    });

    // Read responses from WS (client → relay) and resolve pending requests
    while let Some(Ok(msg)) = ws_stream.next().await {
        if let Message::Text(text) = msg
            && let Ok(resp) = serde_json::from_str::<TunnelResponse>(&text)
        {
            let mut p = pending.write().await;
            if let Some(sender) = p.remove(&resp.id) {
                let _ = sender.send(resp);
            }
        }
    }

    // Client disconnected — clean up
    send_task.abort();
    state.tunnels.remove(&subdomain);
    println!(
        "  {} tunnel disconnected: {}",
        colored::Colorize::red("↓"),
        subdomain
    );
}

/// Catch-all handler: extract subdomain from Host header, forward request through tunnel.
async fn handle_tunnel_request(
    State(state): State<Arc<RelayState>>,
    method: Method,
    headers: HeaderMap,
    uri: axum::http::Uri,
    body: Bytes,
) -> Response {
    // Extract subdomain from Host header: "abc123.tunnel.mcpr.app" → "abc123"
    let host = headers
        .get("host")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");
    let subdomain = host.split('.').next().unwrap_or("").to_string();
    let path_str = uri
        .path_and_query()
        .map(|p| p.to_string())
        .unwrap_or_else(|| "/".into());

    if subdomain.is_empty() {
        relay_log(
            "-",
            method.as_str(),
            &path_str,
            400,
            0,
            std::time::Duration::ZERO,
        );
        return (StatusCode::BAD_REQUEST, "missing host header").into_response();
    }

    // Find tunnel connection
    let conn = match state.tunnels.get(&subdomain) {
        Some(c) => c.clone(),
        None => {
            relay_log(
                &subdomain,
                method.as_str(),
                &path_str,
                502,
                0,
                std::time::Duration::ZERO,
            );
            return (StatusCode::BAD_GATEWAY, "tunnel not found").into_response();
        }
    };

    // Build tunnel request
    let req_id = uuid::Uuid::new_v4().to_string();
    let mut req_headers = HashMap::new();
    for (key, val) in headers.iter() {
        if is_hop_by_hop(key.as_str()) {
            continue;
        }
        if let Ok(v) = val.to_str() {
            req_headers.insert(key.to_string(), v.to_string());
        }
    }

    let body_b64 = if body.is_empty() {
        None
    } else {
        Some(base64::engine::general_purpose::STANDARD.encode(&body))
    };

    let tunnel_req = TunnelRequest {
        id: req_id.clone(),
        method: method.to_string(),
        path: uri
            .path_and_query()
            .map(|p| p.to_string())
            .unwrap_or_else(|| "/".into()),
        headers: req_headers,
        body: body_b64,
    };

    // Register pending response
    let (resp_tx, resp_rx) = oneshot::channel();
    conn.pending.write().await.insert(req_id.clone(), resp_tx);

    // Send request to tunnel client (bounded channel — immediate backpressure)
    let msg = serde_json::to_string(&tunnel_req).unwrap();
    match conn.sender.try_send(msg) {
        Ok(()) => {}
        Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
            conn.pending.write().await.remove(&req_id);
            relay_log(
                &subdomain,
                method.as_str(),
                &path_str,
                503,
                0,
                std::time::Duration::ZERO,
            );
            return (StatusCode::SERVICE_UNAVAILABLE, "tunnel overloaded").into_response();
        }
        Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
            conn.pending.write().await.remove(&req_id);
            relay_log(
                &subdomain,
                method.as_str(),
                &path_str,
                502,
                0,
                std::time::Duration::ZERO,
            );
            return (StatusCode::BAD_GATEWAY, "tunnel disconnected").into_response();
        }
    }

    // Wait for response with timeout
    let path = uri
        .path_and_query()
        .map(|p| p.to_string())
        .unwrap_or_else(|| "/".into());
    let start = std::time::Instant::now();

    match tokio::time::timeout(std::time::Duration::from_secs(30), resp_rx).await {
        Ok(Ok(resp)) => {
            let status_code = StatusCode::from_u16(resp.status).unwrap_or(StatusCode::BAD_GATEWAY);
            let body_len = resp.body.as_ref().map(|b| b.len()).unwrap_or(0);
            relay_log(
                &subdomain,
                method.as_ref(),
                &path,
                resp.status,
                body_len,
                start.elapsed(),
            );

            let mut builder = Response::builder().status(status_code);

            for (k, v) in &resp.headers {
                // Drop hop-by-hop — axum/hyper sets its own framing from the
                // reconstructed body below. Leaving `transfer-encoding` or a
                // stale `content-length` on the response confuses hyper and
                // causes it to abort serialization, closing the TCP
                // connection before any bytes reach the client.
                if is_hop_by_hop(k) {
                    continue;
                }
                if let (Ok(name), Ok(val)) = (
                    HeaderName::from_bytes(k.as_bytes()),
                    HeaderValue::from_str(v),
                ) {
                    builder = builder.header(name, val);
                }
            }

            // Cap response body size: base64 length × 3/4 ≈ decoded size
            let max_b64_len = state.max_response_body * 4 / 3 + 4;
            if resp.body.as_ref().is_some_and(|b| b.len() > max_b64_len) {
                relay_log(
                    &subdomain,
                    method.as_ref(),
                    &path,
                    502,
                    resp.body.as_ref().map(|b| b.len()).unwrap_or(0),
                    start.elapsed(),
                );
                return (StatusCode::BAD_GATEWAY, "tunnel response body too large").into_response();
            }

            let body_bytes = resp
                .body
                .and_then(|b| base64::engine::general_purpose::STANDARD.decode(b).ok())
                .unwrap_or_default();

            builder
                .body(axum::body::Body::from(body_bytes))
                .unwrap_or_else(|_| {
                    (StatusCode::INTERNAL_SERVER_ERROR, "response build error").into_response()
                })
        }
        Ok(Err(_)) => {
            relay_log(&subdomain, method.as_ref(), &path, 502, 0, start.elapsed());
            (StatusCode::BAD_GATEWAY, "tunnel dropped request").into_response()
        }
        Err(_) => {
            conn.pending.write().await.remove(&req_id);
            relay_log(&subdomain, method.as_ref(), &path, 504, 0, start.elapsed());
            (StatusCode::GATEWAY_TIMEOUT, "tunnel timeout").into_response()
        }
    }
}

/// nginx-style access log for relay mode.
fn relay_log(
    subdomain: &str,
    method: &str,
    path: &str,
    status: u16,
    body_len: usize,
    duration: std::time::Duration,
) {
    use colored::Colorize;
    let now = chrono::Local::now().format("%Y-%m-%d %H:%M:%S");
    let status_str = if status < 300 {
        format!("{status}").green().to_string()
    } else if status < 400 {
        format!("{status}").yellow().to_string()
    } else {
        format!("{status}").red().to_string()
    };
    let ms = duration.as_millis();
    println!(
        "  {now}  {sub}  {method} {path}  → {status}  {body_len}b  {ms}ms",
        sub = subdomain.dimmed(),
        status = status_str,
    );
}

// ── Tests ────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(non_snake_case)]
mod tests {
    use super::*;

    #[test]
    fn token_to_subdomain__deterministic() {
        let a = token_to_subdomain("my-token");
        let b = token_to_subdomain("my-token");
        assert_eq!(a, b);
    }

    #[test]
    fn token_to_subdomain__different_tokens_differ() {
        let a = token_to_subdomain("token-a");
        let b = token_to_subdomain("token-b");
        assert_ne!(a, b);
    }

    #[test]
    fn token_to_subdomain__is_12_hex_chars() {
        let sub = token_to_subdomain("test");
        assert_eq!(sub.len(), 12);
        assert!(sub.chars().all(|c| c.is_ascii_hexdigit()));
    }

    #[test]
    fn hex_encode__empty() {
        assert_eq!(hex::encode(&[]), "");
    }

    #[test]
    fn hex_encode__known_values() {
        assert_eq!(hex::encode(&[0x00]), "00");
        assert_eq!(hex::encode(&[0xff]), "ff");
        assert_eq!(hex::encode(&[0xde, 0xad, 0xbe, 0xef]), "deadbeef");
    }

    #[test]
    fn tunnel_request__roundtrip() {
        let req = TunnelRequest {
            id: "req-1".into(),
            method: "POST".into(),
            path: "/mcp".into(),
            headers: HashMap::from([("content-type".into(), "application/json".into())]),
            body: Some(base64::engine::general_purpose::STANDARD.encode(b"{\"test\":true}")),
        };
        let json = serde_json::to_string(&req).unwrap();
        let decoded: TunnelRequest = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.id, "req-1");
        assert_eq!(decoded.method, "POST");
        assert_eq!(decoded.path, "/mcp");
        let body = base64::engine::general_purpose::STANDARD
            .decode(decoded.body.unwrap())
            .unwrap();
        assert_eq!(body, b"{\"test\":true}");
    }

    #[test]
    fn tunnel_response__roundtrip() {
        let resp = TunnelResponse {
            id: "req-1".into(),
            status: 200,
            headers: HashMap::from([("content-type".into(), "application/json".into())]),
            body: None,
        };
        let json = serde_json::to_string(&resp).unwrap();
        let decoded: TunnelResponse = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.id, "req-1");
        assert_eq!(decoded.status, 200);
        assert!(decoded.body.is_none());
    }

    #[test]
    fn register_ack__roundtrip() {
        let ack = RegisterAck {
            subdomain: "abc123".into(),
            url: "https://abc123.tunnel.example.com".into(),
        };
        let json = serde_json::to_string(&ack).unwrap();
        let decoded: RegisterAck = serde_json::from_str(&json).unwrap();
        assert_eq!(decoded.subdomain, "abc123");
        assert_eq!(decoded.url, "https://abc123.tunnel.example.com");
    }

    #[test]
    fn tunnel_request__no_body() {
        let req = TunnelRequest {
            id: "req-2".into(),
            method: "GET".into(),
            path: "/mcp".into(),
            headers: HashMap::new(),
            body: None,
        };
        let json = serde_json::to_string(&req).unwrap();
        let decoded: TunnelRequest = serde_json::from_str(&json).unwrap();
        assert!(decoded.body.is_none());
    }

    // ── build_relay_app ───────────────────────────────────────────────

    fn test_relay_config() -> RelayConfig {
        RelayConfig {
            port: 0, // OS picks
            relay_domain: "tunnel.test".into(),
            auth_provider: None,
            auth_provider_secret: None,
            tokens: HashMap::new(),
            max_request_body_size: None,
            max_response_body_size: None,
        }
    }

    #[test]
    fn build_relay_app__returns_router_and_port() {
        let cfg = test_relay_config();
        let (_, port) = build_relay_app(cfg);
        assert_eq!(port, 0);
    }

    #[test]
    fn build_relay_app__with_static_tokens() {
        let mut tokens = HashMap::new();
        tokens.insert("tok_abc".to_string(), vec!["myapp".to_string()]);
        let cfg = RelayConfig {
            port: 9090,
            relay_domain: "tunnel.test".into(),
            auth_provider: None,
            auth_provider_secret: None,
            tokens,
            max_request_body_size: None,
            max_response_body_size: None,
        };
        let (_, port) = build_relay_app(cfg);
        assert_eq!(port, 9090);
    }

    #[test]
    fn build_relay_app__with_body_size_limits() {
        let cfg = RelayConfig {
            port: 8080,
            relay_domain: "tunnel.test".into(),
            auth_provider: None,
            auth_provider_secret: None,
            tokens: HashMap::new(),
            max_request_body_size: Some(1024),
            max_response_body_size: Some(2048),
        };
        let (_, port) = build_relay_app(cfg);
        assert_eq!(port, 8080);
    }

    #[tokio::test]
    async fn build_relay_app__serves_tunnel_register_endpoint() {
        let cfg = test_relay_config();
        let (app, _) = build_relay_app(cfg);

        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();
        });

        // A GET to /_tunnel/register without WS upgrade headers.
        let resp = reqwest::get(format!("http://{addr}/_tunnel/register"))
            .await
            .unwrap();
        // Not a 404 — route exists (likely 400 since no WS upgrade).
        assert_ne!(resp.status().as_u16(), 404);
    }

    #[tokio::test]
    async fn build_relay_app__fallback_returns_502_for_unknown_subdomain() {
        let cfg = test_relay_config();
        let (app, _) = build_relay_app(cfg);

        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();
        });

        // With no tunnel registered, requests should return 502.
        let resp = reqwest::Client::new()
            .get(format!("http://{addr}/some-path"))
            .header("host", "unknown.tunnel.test")
            .send()
            .await
            .unwrap();
        assert_eq!(resp.status().as_u16(), 502);
    }
}
