// 「網路流量改走你的電腦」: the bot's computer can send its new connections out through this Mac
// (the original: 可以改走你的本機電腦，而不是雲端 — 「僅對新連線生效」, and an admin can turn it off;
// it runs an egress tunnel, Ap.egressConnectProxy 8791). Here: a small HTTP/HTTPS proxy on this
// Mac that the box's commands use through HTTP(S)_PROXY.

use std::io::{Read, Write};
use std::net::{Shutdown, TcpListener, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicU16, AtomicU64, Ordering};
use std::time::Duration;

/// The original's egress CONNECT proxy port; a free one nearby is used when it's taken.
pub const PORT: u16 = 8791;
const HEAD_MAX: usize = 64 * 1024;

static RUNNING: AtomicBool = AtomicBool::new(false);
static BOUND: AtomicU16 = AtomicU16::new(0);
pub static CONNECTIONS: AtomicU64 = AtomicU64::new(0);

/// Start the proxy once (on 127.0.0.1). Returns its port.
pub fn start() -> Result<u16, String> {
    if RUNNING.load(Ordering::SeqCst) {
        return Ok(BOUND.load(Ordering::SeqCst));
    }
    // macOS (Docker Desktop / OrbStack) forwards host.docker.internal to this machine's
    // 127.0.0.1. Linux Docker reaches the host through the bridge (172.17.0.1…), so listen on all
    // addresses there and only serve this machine and its containers (see `allowed_peer`).
    let host = if cfg!(target_os = "linux") { "0.0.0.0" } else { "127.0.0.1" };
    let listener = (PORT..PORT + 20).find_map(|p| TcpListener::bind((host, p)).ok()).ok_or("找不到可用的連接埠")?;
    let port = listener.local_addr().map_err(|e| e.to_string())?.port();
    BOUND.store(port, Ordering::SeqCst);
    RUNNING.store(true, Ordering::SeqCst);
    std::thread::spawn(move || {
        for conn in listener.incoming().flatten() {
            if !conn.peer_addr().map(|a| allowed_peer(a.ip())).unwrap_or(false) {
                let _ = conn.shutdown(Shutdown::Both);
                continue;
            }
            std::thread::spawn(move || {
                let _ = handle(conn);
            });
        }
    });
    Ok(port)
}

/// Who may use the proxy: this machine, and Docker's own networks (172.16.0.0/12, where the
/// bridge and compose networks live) — never the rest of the local network.
pub fn allowed_peer(ip: std::net::IpAddr) -> bool {
    match ip {
        std::net::IpAddr::V4(v) => v.is_loopback() || (v.octets()[0] == 172 && (16..=31).contains(&v.octets()[1])),
        std::net::IpAddr::V6(v) => v.is_loopback(),
    }
}

pub fn port() -> Option<u16> {
    RUNNING.load(Ordering::SeqCst).then(|| BOUND.load(Ordering::SeqCst))
}

/// Where the request goes and what to send there first.
#[derive(Debug, PartialEq)]
pub struct Route {
    pub host: String,
    pub port: u16,
    /// Bytes to send to the server first (empty for CONNECT).
    pub preamble: Vec<u8>,
    pub tunnel: bool,
}

/// Parse a proxy request head: `CONNECT host:443` or `GET http://host[:port]/path HTTP/1.1`.
pub fn route(head: &[u8]) -> Result<Route, String> {
    let text = String::from_utf8_lossy(head);
    let first = text.lines().next().unwrap_or("");
    let mut parts = first.split_whitespace();
    let (method, target, version) = (parts.next().unwrap_or(""), parts.next().unwrap_or(""), parts.next().unwrap_or("HTTP/1.1"));
    if method.eq_ignore_ascii_case("CONNECT") {
        let (h, p) = target.rsplit_once(':').ok_or("CONNECT 要有連接埠")?;
        return Ok(Route { host: h.trim_matches(['[', ']']).to_string(), port: p.parse().map_err(|_| "連接埠不對")?, preamble: vec![], tunnel: true });
    }
    let rest = target.strip_prefix("http://").ok_or("只支援 http:// 或 CONNECT")?;
    let (authority, path) = match rest.find('/') {
        Some(i) => (&rest[..i], &rest[i..]),
        None => (rest, "/"),
    };
    let (host, port) = match authority.rsplit_once(':') {
        Some((h, p)) if p.chars().all(|c| c.is_ascii_digit()) => (h.to_string(), p.parse().unwrap_or(80)),
        _ => (authority.to_string(), 80),
    };
    let mut pre = format!("{method} {path} {version}\r\n").into_bytes();
    let body_start = text.find("\r\n").map(|i| i + 2).unwrap_or(text.len());
    // Drop hop-by-hop proxy headers.
    for line in text[body_start..].split_inclusive("\r\n") {
        let l = line.to_ascii_lowercase();
        if l.starts_with("proxy-connection:") || l.starts_with("proxy-authorization:") {
            continue;
        }
        pre.extend_from_slice(line.as_bytes());
    }
    Ok(Route { host, port, preamble: pre, tunnel: false })
}

fn handle(mut client: TcpStream) -> Result<(), String> {
    client.set_read_timeout(Some(Duration::from_secs(30))).ok();
    let mut buf = vec![];
    let mut chunk = [0u8; 4096];
    let head_end = loop {
        let n = client.read(&mut chunk).map_err(|e| e.to_string())?;
        if n == 0 {
            return Ok(());
        }
        buf.extend_from_slice(&chunk[..n]);
        if let Some(i) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            break i + 4;
        }
        if buf.len() > HEAD_MAX {
            return Err("標頭太大".into());
        }
    };
    let r = match route(&buf[..head_end]) {
        Ok(r) => r,
        Err(e) => {
            let _ = client.write_all(format!("HTTP/1.1 400 Bad Request\r\nContent-Length: {}\r\n\r\n{e}", e.len()).as_bytes());
            return Err(e);
        }
    };
    let mut server = match TcpStream::connect((r.host.as_str(), r.port)) {
        Ok(s) => s,
        Err(e) => {
            let _ = client.write_all(b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\n\r\n");
            return Err(e.to_string());
        }
    };
    CONNECTIONS.fetch_add(1, Ordering::Relaxed);
    client.set_read_timeout(None).ok();
    if r.tunnel {
        client.write_all(b"HTTP/1.1 200 Connection established\r\n\r\n").map_err(|e| e.to_string())?;
    } else {
        server.write_all(&r.preamble).map_err(|e| e.to_string())?;
    }
    // Anything the client sent after the head.
    if head_end < buf.len() {
        server.write_all(&buf[head_end..]).map_err(|e| e.to_string())?;
    }
    let (mut c2, mut s2) = (client.try_clone().map_err(|e| e.to_string())?, server.try_clone().map_err(|e| e.to_string())?);
    let up = std::thread::spawn(move || {
        let _ = std::io::copy(&mut c2, &mut s2);
        let _ = s2.shutdown(Shutdown::Write);
    });
    let _ = std::io::copy(&mut server, &mut client);
    let _ = client.shutdown(Shutdown::Write);
    let _ = up.join();
    Ok(())
}

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

    #[test]
    fn routes_connect_and_plain_http() {
        let r = route(b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n").unwrap();
        assert_eq!((r.host.as_str(), r.port, r.tunnel), ("example.com", 443, true));
        let r = route(b"GET http://example.com:8080/a?b=1 HTTP/1.1\r\nHost: example.com\r\nProxy-Connection: keep-alive\r\n\r\n").unwrap();
        assert_eq!((r.host.as_str(), r.port, r.tunnel), ("example.com", 8080, false));
        let pre = String::from_utf8(r.preamble).unwrap();
        assert!(pre.starts_with("GET /a?b=1 HTTP/1.1\r\n") && !pre.to_lowercase().contains("proxy-connection"));
        assert!(route(b"GET /x HTTP/1.1\r\n\r\n").is_err());
    }

    #[test]
    fn proxies_a_real_request() {
        // A tiny origin server on this machine.
        let origin = TcpListener::bind("127.0.0.1:0").unwrap();
        let op = origin.local_addr().unwrap().port();
        std::thread::spawn(move || {
            if let Ok((mut s, _)) = origin.accept() {
                let mut b = [0u8; 1024];
                let _ = s.read(&mut b);
                let _ = s.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nConnection: close\r\n\r\nhello");
            }
        });
        let port = start().unwrap();
        let mut c = TcpStream::connect(("127.0.0.1", port)).unwrap();
        c.write_all(format!("GET http://127.0.0.1:{op}/ HTTP/1.1\r\nHost: x\r\n\r\n").as_bytes()).unwrap();
        let mut out = String::new();
        c.read_to_string(&mut out).unwrap();
        assert!(out.ends_with("hello"), "{out}");
    }

    #[test]
    fn only_this_machine_and_containers() {
        assert!(super::allowed_peer("127.0.0.1".parse().unwrap()));
        assert!(super::allowed_peer("172.17.0.2".parse().unwrap()));
        assert!(super::allowed_peer("172.31.255.1".parse().unwrap()));
        assert!(!super::allowed_peer("172.32.0.1".parse().unwrap()));
        assert!(!super::allowed_peer("192.168.1.20".parse().unwrap()));
    }
}