// SPEC: PJ2026-01060509 出站诊断 draft-2026-06-26-p8-egress-job-friction.
// Backend-core side of provider egress TCP bridge with stage logs.
use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::bail;
use axum::extract::ws::Message;
use base64::Engine;
use serde_json::{json, Value};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::mpsc;

use crate::state::{AppState, ProviderConnection};

pub struct EgressTcpConnection {
    pub provider_id: String,
    pub connection_id: String,
    pub target_host: String,
    pub target_port: i64,
    pub created_at: Instant,
    pub writer: mpsc::UnboundedSender<Vec<u8>>,
}

const EGRESS_TCP_READ_CHUNK_BYTES: usize = 16 * 1024;
const EGRESS_TCP_COOPERATIVE_YIELD_BYTES: u64 = 128 * 1024;
const EGRESS_TCP_COOPERATIVE_YIELD_DELAY_MS: u64 = 1;

fn egress_tcp_key(provider_id: &str, connection_id: &str) -> String {
    format!("{provider_id}:{connection_id}")
}

fn is_valid_egress_host(host: &str) -> bool {
    !host.is_empty()
        && host.len() <= 253
        && !host
            .chars()
            .any(|ch| ch.is_whitespace() || matches!(ch, '/' | '\\' | ':' | '@' | '\0'))
}

fn is_valid_egress_port(port: i64) -> bool {
    port > 0 && port <= 65_535
}

fn send_egress_close(provider: &Arc<ProviderConnection>, connection_id: &str, error: Option<&str>) {
    let message = match error {
        Some(error) => {
            json!({ "type": "egress_tcp_close", "connectionId": connection_id, "error": error })
        }
        None => json!({ "type": "egress_tcp_close", "connectionId": connection_id }),
    };
    let _ = provider.sender.send(Message::Text(message.to_string()));
}

fn egress_failure_kind(error: &str) -> &'static str {
    let normalized = error.to_ascii_lowercase();
    if normalized.contains("connect timeout") {
        "connect-timeout"
    } else if normalized.contains("timed out") || normalized.contains("timeout") {
        "timeout"
    } else if normalized.contains("invalid") {
        "invalid-target"
    } else if normalized.contains("refused") {
        "connect-refused"
    } else if normalized.contains("unreachable") {
        "connect-unreachable"
    } else {
        "connect-failed"
    }
}

pub async fn handle_egress_tcp_open(
    state: &Arc<AppState>,
    provider: Arc<ProviderConnection>,
    message: &Value,
) -> anyhow::Result<()> {
    let provider_id = message
        .get("providerId")
        .and_then(Value::as_str)
        .unwrap_or("")
        .to_string();
    let connection_id = message
        .get("connectionId")
        .and_then(Value::as_str)
        .unwrap_or("")
        .to_string();
    let host = message
        .get("host")
        .and_then(Value::as_str)
        .unwrap_or("")
        .trim()
        .to_string();
    let port = message.get("port").and_then(Value::as_i64).unwrap_or(0);
    if !is_valid_egress_host(&host)
        || !is_valid_egress_port(port)
        || provider_id.is_empty()
        || connection_id.is_empty()
    {
        state.log(
            "warn",
            "egress_tcp_open_rejected",
            Some(json!({
                "providerId": provider_id.as_str(),
                "connectionId": connection_id.as_str(),
                "targetHost": host.as_str(),
                "targetPort": port,
                "failureKind": "invalid-target"
            })),
        );
        send_egress_close(&provider, &connection_id, Some("invalid egress target"));
        return Ok(());
    }
    state.log(
        "info",
        "egress_tcp_open_requested",
        Some(json!({
            "providerId": provider_id.as_str(),
            "connectionId": connection_id.as_str(),
            "targetHost": host.as_str(),
            "targetPort": port
        })),
    );
    close_egress_tcp_connection(state, &provider_id, &connection_id, None).await;
    state.log(
        "info",
        "egress_tcp_connect_start",
        Some(json!({
            "providerId": provider_id.as_str(),
            "connectionId": connection_id.as_str(),
            "targetHost": host.as_str(),
            "targetPort": port,
            "timeoutMs": 15_000
        })),
    );
    let stream = match tokio::time::timeout(
        Duration::from_secs(15),
        TcpStream::connect((host.as_str(), port as u16)),
    )
    .await
    {
        Ok(Ok(stream)) => stream,
        Ok(Err(error)) => {
            let error_text = error.to_string();
            state.log(
                "warn",
                "egress_tcp_connect_failed",
                Some(json!({
                    "providerId": provider_id.as_str(),
                    "connectionId": connection_id.as_str(),
                    "targetHost": host.as_str(),
                    "targetPort": port,
                    "failureKind": egress_failure_kind(&error_text),
                    "error": error_text.as_str()
                })),
            );
            send_egress_close(&provider, &connection_id, Some(&error_text));
            return Ok(());
        }
        Err(_) => {
            state.log(
                "warn",
                "egress_tcp_connect_timeout",
                Some(json!({
                    "providerId": provider_id.as_str(),
                    "connectionId": connection_id.as_str(),
                    "targetHost": host.as_str(),
                    "targetPort": port,
                    "failureKind": "connect-timeout",
                    "timeoutMs": 15_000
                })),
            );
            send_egress_close(
                &provider,
                &connection_id,
                Some("egress tcp connect timeout"),
            );
            return Ok(());
        }
    };
    let (mut reader, mut writer) = stream.into_split();
    let (tx, mut rx) = mpsc::unbounded_channel::<Vec<u8>>();
    let key = egress_tcp_key(&provider_id, &connection_id);
    state.active_egress_tcp.lock().await.insert(
        key.clone(),
        EgressTcpConnection {
            provider_id: provider_id.clone(),
            connection_id: connection_id.clone(),
            target_host: host.clone(),
            target_port: port,
            created_at: Instant::now(),
            writer: tx,
        },
    );
    let _ = provider.sender.send(Message::Text(
        json!({ "type": "egress_tcp_opened", "connectionId": connection_id.as_str() }).to_string(),
    ));
    state.log(
        "info",
        "egress_tcp_connected",
        Some(json!({
            "providerId": provider_id.as_str(),
            "connectionId": connection_id.as_str(),
            "targetHost": host.as_str(),
            "targetPort": port
        })),
    );

    tokio::spawn({
        let state = state.clone();
        let provider = provider.clone();
        let provider_id = provider_id.clone();
        let connection_id = connection_id.clone();
        let host = host.clone();
        async move {
            let mut buffer = vec![0_u8; EGRESS_TCP_READ_CHUNK_BYTES];
            let mut bytes_to_provider: u64 = 0;
            let mut bytes_since_yield: u64 = 0;
            let mut chunks_to_provider: u64 = 0;
            loop {
                match reader.read(&mut buffer).await {
                    Ok(0) => break,
                    Ok(size) => {
                        let data =
                            base64::engine::general_purpose::STANDARD.encode(&buffer[..size]);
                        let sent = provider.sender.send(Message::Text(
                            json!({
                                "type": "egress_tcp_data",
                                "connectionId": connection_id.as_str(),
                                "data": data,
                                "encoding": "base64"
                            })
                            .to_string(),
                        ));
                        if sent.is_err() {
                            state.log(
                                "warn",
                                "egress_tcp_provider_channel_closed",
                                Some(json!({
                                    "providerId": provider_id.as_str(),
                                    "connectionId": connection_id.as_str(),
                                    "targetHost": host.as_str(),
                                    "targetPort": port,
                                    "failureKind": "provider-channel-closed",
                                    "bytesToProvider": bytes_to_provider,
                                    "chunksToProvider": chunks_to_provider
                                })),
                            );
                            break;
                        }
                        bytes_to_provider += size as u64;
                        bytes_since_yield += size as u64;
                        chunks_to_provider += 1;
                        if bytes_since_yield >= EGRESS_TCP_COOPERATIVE_YIELD_BYTES {
                            bytes_since_yield = 0;
                            tokio::task::yield_now().await;
                            tokio::time::sleep(Duration::from_millis(EGRESS_TCP_COOPERATIVE_YIELD_DELAY_MS)).await;
                        }
                    }
                    Err(error) => {
                        let error_text = error.to_string();
                        state.log(
                            "warn",
                            "egress_tcp_read_failed",
                            Some(json!({
                                "providerId": provider_id.as_str(),
                                "connectionId": connection_id.as_str(),
                                "failureKind": "read-failed",
                                "error": error_text.as_str()
                            })),
                        );
                        send_egress_close(&provider, &connection_id, Some(&error_text));
                        break;
                    }
                }
            }
            state
                .active_egress_tcp
                .lock()
                .await
                .remove(&egress_tcp_key(&provider_id, &connection_id));
            state.log(
                "info",
                "egress_tcp_read_loop_closed",
                Some(json!({
                    "providerId": provider_id.as_str(),
                    "connectionId": connection_id.as_str(),
                    "targetHost": host.as_str(),
                    "targetPort": port,
                    "bytesToProvider": bytes_to_provider,
                    "chunksToProvider": chunks_to_provider,
                    "yieldEveryBytes": EGRESS_TCP_COOPERATIVE_YIELD_BYTES
                })),
            );
            send_egress_close(&provider, &connection_id, None);
        }
    });
    tokio::spawn({
        let state = state.clone();
        let provider_id = provider_id.clone();
        let connection_id = connection_id.clone();
        let host = host.clone();
        async move {
            let mut bytes_from_provider: u64 = 0;
            let mut chunks_from_provider: u64 = 0;
            while let Some(data) = rx.recv().await {
                bytes_from_provider += data.len() as u64;
                chunks_from_provider += 1;
                if let Err(error) = writer.write_all(&data).await {
                    state.log(
                        "warn",
                        "egress_tcp_write_failed",
                        Some(json!({
                            "providerId": provider_id.as_str(),
                            "connectionId": connection_id.as_str(),
                            "failureKind": "write-failed",
                            "error": error.to_string()
                        })),
                    );
                    break;
                }
            }
            state
                .active_egress_tcp
                .lock()
                .await
                .remove(&egress_tcp_key(&provider_id, &connection_id));
            state.log(
                "info",
                "egress_tcp_write_loop_closed",
                Some(json!({
                    "providerId": provider_id.as_str(),
                    "connectionId": connection_id.as_str(),
                    "targetHost": host.as_str(),
                    "targetPort": port,
                    "bytesFromProvider": bytes_from_provider,
                    "chunksFromProvider": chunks_from_provider
                })),
            );
        }
    });
    Ok(())
}

pub async fn handle_egress_tcp_data(state: &Arc<AppState>, message: &Value) -> anyhow::Result<()> {
    let provider_id = message
        .get("providerId")
        .and_then(Value::as_str)
        .unwrap_or("");
    let connection_id = message
        .get("connectionId")
        .and_then(Value::as_str)
        .unwrap_or("");
    let data = message.get("data").and_then(Value::as_str).unwrap_or("");
    let encoding = message
        .get("encoding")
        .and_then(Value::as_str)
        .unwrap_or("utf8");
    let bytes = if encoding == "base64" {
        base64::engine::general_purpose::STANDARD.decode(data)?
    } else {
        data.as_bytes().to_vec()
    };
    let connections = state.active_egress_tcp.lock().await;
    if let Some(connection) = connections.get(&egress_tcp_key(provider_id, connection_id)) {
        let _ = connection.writer.send(bytes);
    }
    Ok(())
}

pub async fn handle_egress_tcp_close(state: &Arc<AppState>, message: &Value) -> anyhow::Result<()> {
    let provider_id = message
        .get("providerId")
        .and_then(Value::as_str)
        .unwrap_or("");
    let connection_id = message
        .get("connectionId")
        .and_then(Value::as_str)
        .unwrap_or("");
    if provider_id.is_empty() || connection_id.is_empty() {
        bail!("egress_tcp_close requires providerId and connectionId");
    }
    close_egress_tcp_connection(state, provider_id, connection_id, None).await;
    Ok(())
}

pub async fn close_egress_tcp_connection(
    state: &Arc<AppState>,
    provider_id: &str,
    connection_id: &str,
    error: Option<&str>,
) {
    let removed_connection = state
        .active_egress_tcp
        .lock()
        .await
        .remove(&egress_tcp_key(provider_id, connection_id));
    let removed = removed_connection.is_some();
    if removed || error.is_some() {
        state.log(
            if error.is_some() { "warn" } else { "info" },
            "egress_tcp_connection_closed",
            Some(json!({
                "providerId": provider_id,
                "connectionId": connection_id,
                "removed": removed,
                "targetHost": removed_connection.as_ref().map(|connection| connection.target_host.as_str()),
                "targetPort": removed_connection.as_ref().map(|connection| connection.target_port),
                "ageMs": removed_connection.as_ref().map(|connection| connection.created_at.elapsed().as_millis() as u64),
                "failureKind": error.map(egress_failure_kind).unwrap_or("normal-close"),
                "error": error.unwrap_or("")
            })),
        );
    }
}

pub async fn close_egress_tcp_connections_for_provider(state: &Arc<AppState>, provider_id: &str) {
    let mut connections = state.active_egress_tcp.lock().await;
    let before = connections.len();
    connections.retain(|_, connection| connection.provider_id != provider_id);
    let removed = before.saturating_sub(connections.len());
    if removed > 0 {
        state.log(
            "warn",
            "egress_tcp_provider_connections_drained",
            Some(json!({
                "providerId": provider_id,
                "removed": removed,
                "failureKind": "provider-disconnected"
            })),
        );
    }
}
