use std::sync::Arc;
use std::time::Duration;

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 writer: mpsc::UnboundedSender<Vec<u8>>,
}

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

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()
    {
        send_egress_close(&provider, &connection_id, Some("invalid egress target"));
        return Ok(());
    }
    close_egress_tcp_connection(state, &provider_id, &connection_id, None).await;
    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)) => {
            send_egress_close(&provider, &connection_id, Some(&error.to_string()));
            return Ok(());
        }
        Err(_) => {
            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(),
            writer: tx,
        },
    );
    let _ = provider.sender.send(Message::Text(
        json!({ "type": "egress_tcp_opened", "connectionId": connection_id }).to_string(),
    ));

    tokio::spawn({
        let state = state.clone();
        let provider = provider.clone();
        let provider_id = provider_id.clone();
        let connection_id = connection_id.clone();
        async move {
            let mut buffer = vec![0_u8; 16 * 1024];
            loop {
                match reader.read(&mut buffer).await {
                    Ok(0) => break,
                    Ok(size) => {
                        let data =
                            base64::engine::general_purpose::STANDARD.encode(&buffer[..size]);
                        let _ = provider.sender.send(Message::Text(json!({ "type": "egress_tcp_data", "connectionId": connection_id, "data": data, "encoding": "base64" }).to_string()));
                    }
                    Err(error) => {
                        send_egress_close(&provider, &connection_id, Some(&error.to_string()));
                        break;
                    }
                }
            }
            state
                .active_egress_tcp
                .lock()
                .await
                .remove(&egress_tcp_key(&provider_id, &connection_id));
            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();
        async move {
            while let Some(data) = rx.recv().await {
                if writer.write_all(&data).await.is_err() {
                    break;
                }
            }
            state
                .active_egress_tcp
                .lock()
                .await
                .remove(&egress_tcp_key(&provider_id, &connection_id));
        }
    });
    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>,
) {
    state
        .active_egress_tcp
        .lock()
        .await
        .remove(&egress_tcp_key(provider_id, connection_id));
}

pub async fn close_egress_tcp_connections_for_provider(state: &Arc<AppState>, provider_id: &str) {
    let mut connections = state.active_egress_tcp.lock().await;
    connections.retain(|_, connection| connection.provider_id != provider_id);
}
