//! `gatunad` — raw-Ethernet tunnel server.
//!
//! Listens on an L2 interface for ethertype `0x6969` frames, announces TCP
//! upstreams to a discovering client, and relays bidirectional TCP traffic
//! between the client and `127.0.0.1:<port>` services.

mod frame;
mod link;
mod session;
mod upstream;

use clap::Parser;
use std::process::ExitCode;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, Mutex};
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio::sync::Mutex as AsyncMutex;
use tracing::{error, info, Level};

use crate::frame::{
    Frame, REASON_CONNECT_FAILED, REASON_UNKNOWN_SESSION, REASON_UNKNOWN_UPSTREAM,
};
use crate::link::Link;
use crate::session::{spawn_pump, handle_data, send_pure_ack, SessionHandle, SessionStore};
use crate::upstream::build_table;

#[derive(Parser)]
#[command(name = "gatunad", version, about = "raw ethernet tunnel server")]
struct Args {
    /// L2 interface name (e.g. eth0).
    iface: String,
    /// One or more TCP upstreams as PORT[:label], relayed to 127.0.0.1:PORT.
    #[arg(num_args = 1..)]
    ports: Vec<String>,
    /// Verbose logging (lifecycle events to stdout).
    #[arg(short, long)]
    verbose: bool,
}

#[tokio::main]
async fn main() -> ExitCode {
    let level = if args.verbose { Level::INFO } else { Level::ERROR };
    tracing_subscriber::fmt()
        .with_max_level(level)
        .with_writer(|| std::io::stdout())
        .init();

    let args = Args::parse();

    let table = match build_table(&args.ports) {
        Ok(t) => Arc::new(t),
        Err(e) => {
            eprintln!("gatunad: {e}");
            return ExitCode::FAILURE;
        }
    };

    let hostname = Arc::new(upstream::get_hostname());

    let link = match Link::open(&args.iface) {
        Ok(l) => Arc::new(l),
        Err(e) => {
            eprintln!("gatunad: failed to open interface {}: {e}", args.iface);
            return ExitCode::FAILURE;
        }
    };

    let (tx, mut rx) = mpsc::channel::<(crate::link::MacAddr, Vec<u8>)>(1024);

    // Tx task: sole owner of send-side writes, fed by all session pumps + rx.
    {
        let link = Arc::clone(&link);
        tokio::spawn(async move {
            while let Some((dst, frame)) = rx.recv().await {
                loop {
                    let mut guard = match link.writable().await {
                        Ok(g) => g,
                        Err(e) => {
                            error!("writable wait: {e}");
                            break;
                        }
                    };
                    match link.send(dst, &frame) {
                        Ok(()) => break,
                        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                            guard.clear_ready();
                            continue;
                        }
                        Err(e) => {
                            error!("send failed: {e}");
                            break;
                        }
                    }
                }
            }
        });
    }

    let store: SessionStore = Arc::new(Mutex::new(std::collections::HashMap::new()));
    let next_id = Arc::new(AtomicU32::new(1));
    let mut buf = vec![0u8; 65536];

    loop {
        let mut guard = match link.readable().await {
            Ok(g) => g,
            Err(e) => {
                error!("readable wait: {e}");
                continue;
            }
        };
        match link.recv(&mut buf) {
            Ok(Some((src, payload))) => {
                let frame = match Frame::parse(payload) {
                    Ok(f) => f,
                    Err(e) => {
                        error!("decode from {src:?}: {e}");
                        continue;
                    }
                };
                handle_frame(frame, src, &tx, &store, &next_id, &table, &hostname).await;
            }
            Ok(None) => {
                // Ignorable frame (outgoing/short/mismatch) or transient; do not
                // clear readiness — keep draining.
            }
            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                guard.clear_ready();
            }
            Err(e) => {
                error!("recv: {e}");
            }
        }
    }
}

async fn handle_frame(
    frame: Frame,
    src: crate::link::MacAddr,
    tx: &mpsc::Sender<(crate::link::MacAddr, Vec<u8>)>,
    store: &SessionStore,
    next_id: &Arc<AtomicU32>,
    table: &Arc<crate::upstream::UpstreamTable>,
    hostname: &Arc<String>,
) {
    match frame {
        Frame::Discover => {
            info!("DISCOVER from {src:?}");
            let manifest = Frame::Manifest {
                hostname: (**hostname).clone(),
                entries: table.entries(),
            };
            let _ = tx.send((src, manifest.encode())).await;
            info!("MANIFEST sent to {src:?} ({} upstreams)", table.0.len());
        }
        Frame::Open { upstream_id, proto: _ } => {
            info!("OPEN upstream {upstream_id} from {src:?}");
            let upstream = table.get(upstream_id);
            match upstream {
                Some(upstream) => {
                    let port = upstream.port;
                    let proto = upstream.proto.as_u8();
                    // Spawn so connect() doesn't block the rx loop.
                    let tx = tx.clone();
                    let store = Arc::clone(store);
                    let next = Arc::clone(next_id);
                    tokio::spawn(async move {
                        match TcpStream::connect(("127.0.0.1", port)).await {
                            Ok(stream) => {
                                let sid = next.fetch_add(1, Ordering::Relaxed);
                                let (r, w) = stream.into_split();
                                let w = Arc::new(AsyncMutex::new(w));
                                store.lock().expect("store poisoned").insert(
                                    sid,
                                    SessionHandle {
                                        upstream_id,
                                        proto,
                                        write: w,
                                        send_state: Arc::new(Mutex::new(session::SendState::new())),
                                        recv_state: Arc::new(Mutex::new(session::RecvState::new())),
                                        peer_mac: src,
                                    },
                                );
                                let ack = Frame::OpenAck { session_id: sid, upstream_id, proto };
                                let _ = tx.send((src, ack.encode())).await;
                                info!("session {sid} established (upstream {upstream_id})");
                                spawn_pump(r, sid, proto, src, tx, store);
                            }
                            Err(e) => {
                                error!("connect 127.0.0.1:{port} failed: {e}");
                                info!("OPEN_NAK upstream {upstream_id} (connect_failed) to {src:?}");
                                let nak =
                                    Frame::OpenNak { upstream_id, reason: REASON_CONNECT_FAILED };
                                let _ = tx.send((src, nak.encode())).await;
                            }
                        }
                    });
                }
                None => {
                    let nak = Frame::OpenNak { upstream_id, reason: REASON_UNKNOWN_UPSTREAM };
                    let _ = tx.send((src, nak.encode())).await;
                }
            }
        }
        Frame::Data { session_id, seq, ack_seq, payload } => {
            match handle_data(store, session_id, seq, ack_seq, &payload, tx).await {
                Ok(need_ack) => {
                    if need_ack {
                        send_pure_ack(store, session_id, tx);
                    }
                }
                Err(session::WriteError::UnknownSession) => {
                    let close = Frame::Close {
                        session_id,
                        reason: Some(REASON_UNKNOWN_SESSION),
                    };
                    let _ = tx.send((src, close.encode())).await;
                }
                Err(session::WriteError::Io) => {
                    store.lock().expect("store poisoned").remove(&session_id);
                    let close = Frame::Close { session_id, reason: None };
                    let _ = tx.send((src, close.encode())).await;
                }
            }
        }
        Frame::Close { session_id, reason: _ } => {
            info!("CLOSE session {session_id} from {src:?}");
            store.lock().expect("store poisoned").remove(&session_id);
        }
        Frame::Ping { nonce } => {
            let pong = Frame::Pong { nonce };
            let _ = tx.send((src, pong.encode())).await;
        }
        // Not expected from a client; ignore.
        Frame::Manifest { .. } | Frame::OpenAck { .. } | Frame::OpenNak { .. }
        | Frame::Pong { .. } => {}
    }
}