use crate::mut_qstring::MutQString;
use crate::tun::{BorrowedTunIo, RewindStream};
use crate::{error, info, warn};
use bytes::Bytes;
use http::Uri;
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper_util::rt::TokioIo;
use ipstack::IpStackStream;
use regex::Regex;
use rustls::RootCertStore;
use std::io;
use std::net::SocketAddr;
use std::os::fd::RawFd;
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::UdpSocket;
use tokio::task::JoinSet;
use tokio_rustls::StartHandshake;
use tokio_util::sync::CancellationToken;
use tun_rs::BorrowedAsyncDevice;
use crate::acceptor::build_tls_acceptor;

static ROOT_STORE: OnceLock<Arc<RootCertStore>> = OnceLock::new();

const BLOCKED_HOSTS: &[&str] = &["submit.backtrace.io"];

fn is_blocked_host(host: &str) -> bool {
    let host = host.trim_end_matches('.');
    BLOCKED_HOSTS
        .iter()
        .any(|blocked| host.eq_ignore_ascii_case(blocked))
}

const SNIFF_TIMEOUT: Duration = Duration::from_secs(5);
const TEARDOWN_DRAIN_TIMEOUT: Duration = Duration::from_secs(3);
const TEARDOWN_SETTLE: Duration = Duration::from_millis(50);
const TEARDOWN_FIN_TIMEOUT: Duration = Duration::from_millis(300);
const TEARDOWN_GRACE: Duration = Duration::from_secs(4);

fn root_store() -> Arc<RootCertStore> {
    ROOT_STORE
        .get_or_init(|| {
            let mut root_store = RootCertStore::empty();
            let native_certs = rustls_native_certs::load_native_certs();
            for err in native_certs.errors {
                warn!("加载证书警告: {err}");
            }
            for cert in native_certs.certs {
                let _ = root_store.add(cert);
            }

            load_system_ca_dirs(&mut root_store);

            if root_store.is_empty() {
                root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
            }
            Arc::new(root_store)
        })
        .clone()
}

fn load_system_ca_dirs(store: &mut RootCertStore) {
    use std::io::BufReader;

    const DIRS: &[&str] = &[
        "/system/etc/security/cacerts",        // Android 13-
        "/apex/com.android.conscrypt/cacerts", // Android 14+
    ];

    for dir in DIRS {
        let Ok(entries) = std::fs::read_dir(dir) else {
            continue;
        };
        for entry in entries.flatten() {
            let Ok(file) = std::fs::File::open(entry.path()) else {
                continue;
            };
            for cert in rustls_pemfile::certs(&mut BufReader::new(file)).flatten() {
                let _ = store.add(cert);
            }
        }
    }
}

async fn sniff_client_hello<S>(
    io: &mut RewindStream<S>,
) -> io::Result<Option<(Option<String>, rustls::server::Accepted)>>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let mut acceptor = rustls::server::Acceptor::default();
    let mut buf = vec![0u8; 16 * 1024];

    loop {
        let n = match tokio::time::timeout(SNIFF_TIMEOUT, io.read(&mut buf)).await {
            Ok(Ok(n)) => n,
            Ok(Err(e)) => return Err(e),
            Err(_) => return Ok(None), // 对端迟迟不发数据，当作非 TLS
        };
        if n == 0 {
            return Ok(None);
        }

        if acceptor.read_tls(&mut &buf[..n]).is_err() {
            return Ok(None);
        }
        match acceptor.accept() {
            Ok(Some(accepted)) => {
                let sni = accepted.client_hello().server_name().map(|s| s.to_string());
                return Ok(Some((sni, accepted)));
            }
            Ok(None) => continue,      // 还没凑齐
            Err(_) => return Ok(None), // 不是合法 TLS
        }
    }
}

#[allow(clippy::too_many_arguments)]
async fn handle_tls_connection<S>(
    tcp: S,
    dst: SocketAddr,
    upstream: Upstream,
    target_hosts: Vec<String>,
    waiters: Waiters,
    server_config: Arc<rustls::ServerConfig>,
    cancel: CancellationToken,
) -> anyhow::Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let mut io = RewindStream::new(tcp);
    let sniffed = match sniff_client_hello(&mut io).await {
        Ok(v) => v,
        Err(e) => {
            warn!("client hello sniff error on {dst:?}: {e:?}");
            return Ok(());
        }
    };

    let Some((sni, accepted)) = sniffed else {
        io.rewind();
        io.stop_recording();
        return passthrough_cancel_wrapped(io, dst, None, upstream, &cancel).await;
    };

    if let Some(host) = sni.as_deref()
        && is_blocked_host(host)
    {
        warn!("blocked host {host} ({dst:?}), connection rejected");
        return Ok(());
    }

    let host_matched = sni
        .as_deref()
        .map(|s| target_hosts.iter().any(|h| s.eq_ignore_ascii_case(h)))
        .unwrap_or(false);

    if !host_matched {
        io.rewind();
        io.stop_recording();
        let r = passthrough_cancel_wrapped(io, dst, sni.as_deref(), upstream, &cancel).await;
        if r.is_err() {
            warn!("tcp_passthrough error on {sni:?} {dst:?}");
        }
        return r;
    }

    io.stop_recording();
    let client_tls = StartHandshake::from_parts(accepted, io)
        .into_stream(server_config)
        .await?;

    let svc_host = sni.clone().unwrap_or_default();
    let svc = hyper::service::service_fn(move |req: hyper::Request<Incoming>| {
        let host = svc_host.clone();
        let waiters = waiters.clone();
        let upstream = upstream.clone();
        async move {
            Ok::<_, std::convert::Infallible>(
                forward_request(req, dst, host, waiters, upstream).await,
            )
        }
    });

    let conn =
        hyper::server::conn::http1::Builder::new().serve_connection(TokioIo::new(client_tls), svc);
    tokio::pin!(conn);

    let mut graceful = false;
    loop {
        tokio::select! {
            res = conn.as_mut() => {
                match res {
                    Ok(()) => {}
                    Err(e) => warn!("downstream http1 connection error: {e:?}"),
                }
                break;
            }
            _ = cancel.cancelled(), if !graceful => {
                info!("teardown: graceful shutdown of downstream connection");
                conn.as_mut().graceful_shutdown();
                graceful = true;
            }
        }
    }

    Ok(())
}

async fn passthrough_cancel_wrapped<IO>(
    io: IO,
    dst: SocketAddr,
    sni: Option<&str>,
    upstream: Upstream,
    cancel: &CancellationToken,
) -> anyhow::Result<()>
where
    IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    tokio::select! {
        _ = cancel.cancelled() => Ok(()),
        r = tcp_passthrough(io, dst, sni, upstream) => r,
    }
}

#[allow(clippy::too_many_arguments)]
async fn forward_request(
    mut req: hyper::Request<Incoming>,
    dst: SocketAddr,
    url_host: String,
    waiters: Waiters,
    upstream: Upstream,
) -> hyper::Response<Full<Bytes>> {
    if let Some(rewritten) = force_full_mysekai(req.uri()) {
        warn!("rewrite mysekai flag: {} -> {}", req.uri(), rewritten);
        *req.uri_mut() = rewritten;
    }

    let uri = req.uri().clone();
    let url = build_url(&url_host, &uri);
    let path_and_query = uri
        .path_and_query()
        .map(|p| p.as_str().to_string())
        .unwrap_or_else(|| "/".to_string());

    let waiter = take_matching_waiter(&waiters, &path_and_query);
    if let Some(w) = &waiter {
        info!("step '{}' matched", w.id);
    }
    let want_body = waiter.is_some();

    match forward_request_inner(req, dst, &url_host, upstream, want_body).await {
        Ok((response, saved_body)) => {
            if let (Some(w), Some(body)) = (waiter, saved_body) {
                save_and_notify(w, url, body);
            }
            response
        }
        Err(e) => {
            error!("forward to {dst} failed: {e:?}");
            if let Some(w) = waiter {
                crate::notify_step_result(&w.id, None);
            }
            hyper::Response::builder()
                .status(hyper::StatusCode::BAD_GATEWAY)
                .header("content-type", "text/plain; charset=utf-8")
                .body(Full::new(Bytes::from(format!("mitm forward error: {e}"))))
                .unwrap()
        }
    }
}

fn save_and_notify(waiter: Waiter, url: String, body: Bytes) {
    tokio::spawn(async move {
        let path = std::path::Path::new(&waiter.save_path);
        if !path.is_absolute() {
            error!("save_path 不是绝对路径（{}）", waiter.save_path);
        }
        if let Some(parent) = path.parent() {
            let _ = tokio::fs::create_dir_all(parent).await;
        }
        match tokio::fs::write(path, &body).await {
            Ok(()) => {
                info!("step '{}' saved ({} bytes)", waiter.id, body.len());
                crate::notify_step_result(&waiter.id, Some(&url));
            }
            Err(e) => {
                error!("write {} failed: {e:?}", waiter.save_path);
                crate::notify_step_result(&waiter.id, None);
            }
        }
    });
}

async fn forward_request_inner(
    req: hyper::Request<Incoming>,
    dst: SocketAddr,
    url_host: &str,
    upstream: Upstream,
    want_body: bool,
) -> anyhow::Result<(hyper::Response<Full<Bytes>>, Option<Bytes>)> {
    let (mut parts, body) = req.into_parts();
    let body_bytes = body.collect().await?.to_bytes();

    // 拒绝长链接
    for h in [
        "connection",
        "keep-alive",
        "proxy-connection",
        "transfer-encoding",
        "upgrade",
        "te",
        "trailer",
        "content-length",
    ] {
        parts.headers.remove(h);
    }

    let mut client_config = rustls::ClientConfig::builder()
        .with_root_certificates(root_store())
        .with_no_client_auth();
    client_config.alpn_protocols = vec![b"http/1.1".to_vec()];
    let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));
    let server_name = rustls::pki_types::ServerName::try_from(url_host.to_string())?;
    let tcp_to_server = upstream.connect(dst, Some(url_host)).await?;
    let server_tls = connector.connect(server_name, tcp_to_server).await?;

    let (mut sender, conn) =
        hyper::client::conn::http1::handshake::<_, Full<Bytes>>(TokioIo::new(server_tls)).await?;
    tokio::spawn(async move {
        if let Err(e) = conn.await {
            warn!("upstream http1 connection error: {e:?}");
        }
    });

    let upstream_req = hyper::Request::from_parts(parts, Full::new(body_bytes));
    let resp = sender.send_request(upstream_req).await?;

    let (mut res_parts, res_body) = resp.into_parts();
    let res_bytes = res_body.collect().await?.to_bytes();
    for h in [
        "connection",
        "keep-alive",
        "proxy-connection",
        "transfer-encoding",
        "upgrade",
        "te",
        "trailer",
    ] {
        res_parts.headers.remove(h);
    }

    let saved = if want_body {
        Some(res_bytes.clone())
    } else {
        None
    };
    Ok((
        hyper::Response::from_parts(res_parts, Full::new(res_bytes)),
        saved,
    ))
}

fn build_url(host: &str, uri: &Uri) -> String {
    if uri.scheme().is_some() && uri.authority().is_some() {
        return uri.to_string();
    }
    let pq = uri.path_and_query().map(|p| p.as_str()).unwrap_or("/");
    format!("https://{host}{pq}")
}

fn force_full_mysekai(uri: &Uri) -> Option<Uri> {
    const KEY: &str = "isForceAllReloadOnlyMysekai";
    let mut query = qstring::QString::from(uri.query()?);
    if query.get(KEY) != Some("False") {
        return None;
    }
    query.set(KEY, "True");
    Uri::try_from(format!("{}?{}", uri.path(), query).as_str()).ok()
}

#[derive(serde::Deserialize, Debug, Clone)]
pub struct StepSpec {
    pub id: String,
    #[serde(default)]
    pub hosts: Vec<String>,
    pub path: String,
    #[serde(rename = "savePath")]
    pub save_path: String,
    /// 持续运行：命中后不摘除 waiter，同类型的包可以反复抓
    #[serde(default)]
    pub repeat: bool,
}

#[derive(Clone)]
pub struct Waiter {
    pub id: String,
    pub path: Regex,
    pub save_path: String,
    pub repeat: bool,
}

pub type Waiters = Arc<StdMutex<Vec<Waiter>>>;

fn take_matching_waiter(waiters: &StdMutex<Vec<Waiter>>, path_and_query: &str) -> Option<Waiter> {
    let mut list = waiters.lock().unwrap();
    let idx = list.iter().position(|w| w.path.is_match(path_and_query))?;
    if list[idx].repeat {
        // 持续运行：保留 waiter，让后续同类型的请求继续命中
        Some(list[idx].clone())
    } else {
        Some(list.remove(idx))
    }
}

#[derive(Clone, Debug)]
pub struct ProxyConfig {
    pub tls: bool,
    pub host: String,
    pub port: u16,
    pub auth: Option<String>,
}

impl ProxyConfig {
    pub fn parse(input: &str) -> anyhow::Result<Option<Self>> {
        let raw = input.trim();
        if raw.is_empty() {
            return Ok(None);
        }
        let url = url::Url::parse(raw)
            .map_err(|e| anyhow::anyhow!("代理地址不是合法 URL: {raw} ({e})"))?;
        let tls = match url.scheme() {
            "http" => false,
            "https" => true,
            other => anyhow::bail!("代理只支持 http/https，当前是 {other}://"),
        };
        let host = url
            .host_str()
            .filter(|h| !h.is_empty())
            .ok_or_else(|| anyhow::anyhow!("代理地址缺少 host: {raw}"))?
            .to_string();
        let port = url
            .port_or_known_default()
            .ok_or_else(|| anyhow::anyhow!("代理地址缺少端口: {raw}"))?;
        let auth = if url.username().is_empty() {
            None
        } else {
            let pass = url.password().unwrap_or("");
            Some(base64_encode(
                format!("{}:{}", url.username(), pass).as_bytes(),
            ))
        };
        Ok(Some(Self {
            tls,
            host,
            port,
            auth,
        }))
    }
}

fn base64_encode(data: &[u8]) -> String {
    use base64::Engine as _;
    base64::engine::general_purpose::STANDARD.encode(data)
}

type UpstreamIo = Box<dyn AsyncReadWrite + Send + Unpin>;

trait AsyncReadWrite: AsyncRead + AsyncWrite {}
impl<T: AsyncRead + AsyncWrite + ?Sized> AsyncReadWrite for T {}

#[derive(Clone)]
struct Upstream {
    proxy: Option<Arc<ProxyConfig>>,
}

impl Upstream {
    fn new(proxy: Option<ProxyConfig>) -> Self {
        Self {
            proxy: proxy.map(Arc::new),
        }
    }

    async fn connect(&self, dst: SocketAddr, host: Option<&str>) -> anyhow::Result<UpstreamIo> {
        match &self.proxy {
            None => {
                let tcp = tokio::net::TcpStream::connect(dst).await?;
                tcp.set_nodelay(true).ok();
                Ok(Box::new(tcp))
            }
            Some(proxy) => connect_via_proxy(proxy, dst, host).await,
        }
    }
}

async fn connect_via_proxy(
    proxy: &ProxyConfig,
    dst: SocketAddr,
    host: Option<&str>,
) -> anyhow::Result<UpstreamIo> {
    let tcp = tokio::net::TcpStream::connect((proxy.host.as_str(), proxy.port))
        .await
        .map_err(|e| anyhow::anyhow!("连接代理 {}:{} 失败: {e}", proxy.host, proxy.port))?;
    tcp.set_nodelay(true).ok();

    let mut io: UpstreamIo = if proxy.tls {
        let mut cfg = rustls::ClientConfig::builder()
            .with_root_certificates(root_store())
            .with_no_client_auth();
        cfg.alpn_protocols = vec![b"http/1.1".to_vec()];
        let connector = tokio_rustls::TlsConnector::from(Arc::new(cfg));
        let name = rustls::pki_types::ServerName::try_from(proxy.host.clone())?;
        Box::new(
            connector
                .connect(name, tcp)
                .await
                .map_err(|e| anyhow::anyhow!("与代理建立 TLS 失败: {e}"))?,
        )
    } else {
        Box::new(tcp)
    };

    let authority = match host {
        Some(h) if !h.is_empty() => format!("{h}:{}", dst.port()),
        _ => dst.to_string(),
    };
    let mut req = format!("CONNECT {authority} HTTP/1.1\r\nHost: {authority}\r\n");
    if let Some(auth) = &proxy.auth {
        req.push_str(&format!("Proxy-Authorization: Basic {auth}\r\n"));
    }
    req.push_str("\r\n");
    io.write_all(req.as_bytes()).await?;
    io.flush().await?;

    let mut head = Vec::with_capacity(256);
    let mut byte = [0u8; 1];
    loop {
        let n = io.read(&mut byte).await?;
        if n == 0 {
            anyhow::bail!("代理在 CONNECT 阶段直接断开");
        }
        head.push(byte[0]);
        if head.ends_with(b"\r\n\r\n") {
            break;
        }
        if head.len() > 16 * 1024 {
            anyhow::bail!("代理 CONNECT 响应头过长");
        }
    }
    let head_str = String::from_utf8_lossy(&head);
    let status_line = head_str.lines().next().unwrap_or("").trim().to_string();
    let ok = status_line
        .split_whitespace()
        .nth(1)
        .map(|code| code == "200")
        .unwrap_or(false);
    if !ok {
        anyhow::bail!("代理 CONNECT 失败: {status_line}");
    }
    Ok(io)
}

pub async fn run_session(
    tun_fd: i32,
    target_hosts: Vec<String>,
    proxy: Option<ProxyConfig>,
    waiters: Waiters,
    cancel: CancellationToken,
) -> anyhow::Result<()> {
    let device = unsafe { BorrowedAsyncDevice::borrow_raw(tun_fd as RawFd)? };

    let mut ipstack_config = ipstack::IpStackConfig::default();
    ipstack_config.mtu(1500)?;
    let device_ref: &BorrowedAsyncDevice<'static> = unsafe { std::mem::transmute(&device) };
    let tun_io = BorrowedTunIo::new(device_ref);
    let mut ip_stack = ipstack::IpStack::new(ipstack_config, tun_io);
    let mut bridge_tasks = JoinSet::new();

    let tls_acceptor = build_tls_acceptor(target_hosts.clone());
    let upstream = Upstream::new(proxy);

    loop {
        let token = cancel.clone();
        tokio::select! {
            _ = cancel.cancelled() => {
                break;
            }
            _ = bridge_tasks.join_next(), if !bridge_tasks.is_empty() => {}

            res = ip_stack.accept() => {
                let stream = match res {
                    Ok(s) => s,
                    Err(e) => {
                        error!("ip_stack accept error (device might be dropped): {e:?}");
                        break;
                    }
                };

                match stream {
                    IpStackStream::Udp(udp) => {
                        let dst = udp.peer_addr();
                        if dst.port() == 443 {
                            // drop quic roughly
                            drop(udp);
                        } else {
                            bridge_tasks.spawn(async move {
                                if let Err(e) = udp_passthrough(udp, dst, token).await {
                                    error!("udp passthrough error: {e:?}");
                                }
                            });
                        }
                    }
                    IpStackStream::Tcp(mut tcp) => {
                        let dst = tcp.peer_addr();
                        if dst.port() == 443 {
                            let (app_io, mut tun_io) = tokio::io::duplex(256 * 1024);

                            let pump_token = cancel.clone();
                            bridge_tasks.spawn(async move {
                                tokio::select! {
                                    _ = pump_token.cancelled() => {
                                        let _ = tokio::time::timeout(
                                            TEARDOWN_DRAIN_TIMEOUT,
                                            tokio::io::copy(&mut tun_io, &mut tcp),
                                        )
                                        .await;
                                    }
                                    _ = tokio::io::copy_bidirectional(&mut tun_io, &mut tcp) => {}
                                }
                                tokio::time::sleep(TEARDOWN_SETTLE).await;
                                let _ = tokio::time::timeout(
                                    TEARDOWN_FIN_TIMEOUT,
                                    tcp.shutdown(),
                                )
                                .await;
                            });

                            let hosts = target_hosts.clone();
                            let server_config = tls_acceptor.clone();
                            let token = cancel.clone();
                            let waiters = waiters.clone();
                            let upstream = upstream.clone();

                            tokio::spawn(async move {
                                match handle_tls_connection(
                                    app_io,
                                    dst,
                                    upstream,
                                    hosts,
                                    waiters,
                                    server_config,
                                    token,
                                ).await {
                                    Ok(()) => {}
                                    Err(e) => error!("handle tls conn error: {e:?}"),
                                }
                            });
                        } else {
                            let pump_token = cancel.clone();
                            let upstream = upstream.clone();
                            bridge_tasks.spawn(async move {
                                let _ =
                                    tcp_passthrough_cancellable(tcp, dst, None, upstream, pump_token)
                                        .await;
                            });
                        }
                    }
                    _ => {}
                }
            }
        }
    }

    let _ = tokio::time::timeout(TEARDOWN_GRACE, async {
        while bridge_tasks.join_next().await.is_some() {}
    })
    .await;
    bridge_tasks.abort_all();
    while bridge_tasks.join_next().await.is_some() {}

    info!("session stopped");
    Ok(())
}

async fn tcp_passthrough<IO>(
    mut client: IO,
    dst: SocketAddr,
    sni: Option<&str>,
    upstream: Upstream,
) -> anyhow::Result<()>
where
    IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let mut server = upstream.connect(dst, sni).await?;
    tokio::io::copy_bidirectional(&mut client, &mut server).await?;
    Ok(())
}

async fn tcp_passthrough_cancellable<IO>(
    mut client: IO,
    dst: SocketAddr,
    sni: Option<&str>,
    upstream: Upstream,
    cancel: CancellationToken,
) -> anyhow::Result<()>
where
    IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let mut server = upstream.connect(dst, sni).await?;
    tokio::select! {
        _ = cancel.cancelled() => {
            let _ = tokio::time::timeout(
                TEARDOWN_DRAIN_TIMEOUT,
                tokio::io::copy(&mut client, &mut server),
            )
            .await;
        }
        res = tokio::io::copy_bidirectional(&mut client, &mut server) => {
            res?;
        }
    }
    tokio::time::sleep(TEARDOWN_SETTLE).await;
    let _ = tokio::time::timeout(TEARDOWN_FIN_TIMEOUT, client.shutdown()).await;
    let _ = tokio::time::timeout(TEARDOWN_FIN_TIMEOUT, server.shutdown()).await;
    Ok(())
}
async fn udp_passthrough<IO>(
    client: IO,
    dst: SocketAddr,
    cancel: CancellationToken,
) -> anyhow::Result<()>
where
    IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
    let bind_addr: SocketAddr = if dst.is_ipv4() {
        "0.0.0.0:0".parse()?
    } else {
        "[::]:0".parse()?
    };

    let socket = UdpSocket::bind(bind_addr).await?;

    socket.connect(dst).await?;
    let socket = Arc::new(socket);

    let (mut cr, mut cw) = tokio::io::split(client);

    let mut client_buf = vec![0u8; 65535];
    let mut server_buf = vec![0u8; 65535];

    let idle_timeout = Duration::from_secs(12);

    loop {
        tokio::select! {
            _ = cancel.cancelled() => break,

            // c2s
            res = tokio::time::timeout(idle_timeout, cr.read(&mut client_buf)) => {
                match res {
                    Ok(Ok(0)) | Err(_) => break,
                    Ok(Ok(n)) => {
                        if socket.send(&client_buf[..n]).await.is_err() {
                            break;
                        }
                    }
                    Ok(Err(_)) => break,
                }
            }

            // s2c
            res = tokio::time::timeout(idle_timeout, socket.recv(&mut server_buf)) => {
                match res {
                    Ok(Ok(n)) if n > 0 => {
                        if cw.write_all(&server_buf[..n]).await.is_err() {
                            break;
                        }
                    }
                    _ => break,
                }
            }
        }
    }

    Ok(())
}