regy 0.1.0

use std::{net::SocketAddr, sync::Arc, time::Duration};

use iroh::{
    Endpoint, RelayMode, Watcher,
    endpoint::{QuicTransportConfig, presets},
};
use iroh_tickets::endpoint::EndpointTicket;
use regy_ui_wire::{FrameLimit, JsonFrameDecoder, encode_json_frame};
use serde_json::Value;
use tokio::{
    io::AsyncWriteExt as _,
    sync::{Semaphore, mpsc, watch},
    task::{JoinHandle, JoinSet},
    time::timeout,
};
use tokio_util::{sync::CancellationToken, task::AbortOnDropHandle};

use crate::{
    config::app::IrohRelayConfig,
    domain::errors::{AgentError, AgentResult, ErrorCode},
    pairing::{ClientBinding, ConnectionId, SourceIdentity},
    presentation::{
        ui_channel::{UI_CLIENT_QUEUE_CAPACITY, UiChannel, UiChannelPeer, UiHub},
        ui_handshake::{UiPairingService, authenticate_authenticated_channel},
    },
    transport::iroh_identity::IrohIdentityStore,
};

pub(crate) const UI_IROH_ALPN: &[u8] = b"regy/ui/1";

const KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(5);
const MAX_IDLE_TIMEOUT: Duration = Duration::from_secs(30);
const STREAM_ACCEPT_TIMEOUT: Duration = Duration::from_secs(5);
const ENDPOINT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(2);
const TERMINAL_HANDSHAKE_DRAIN_TIMEOUT: Duration = Duration::from_secs(1);
const MAX_CONCURRENT_CONNECTIONS: usize = 16;
const MAX_CONCURRENT_BIDIRECTIONAL_STREAMS: u32 = 1;
const MAX_CONCURRENT_UNIDIRECTIONAL_STREAMS: u32 = 0;
const CLOSE_STREAM_TIMEOUT: u32 = 0x101;
const CLOSE_STREAM_LIMIT: u32 = 0x102;
const CLOSE_PROTOCOL: u32 = 0x103;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum IrohEndpointRoute {
    RelayOnly,
    // Retained for loopback integration tests; public startup constructs RelayOnly.
    #[allow(dead_code)]
    DirectLoopback,
}

#[derive(Clone)]
enum IrohRelayRoute {
    N0,
    Explicit(Vec<iroh::RelayUrl>),
    // Retained for the loopback endpoint plan used by integration tests.
    #[allow(dead_code)]
    Disabled,
}

/// A testable representation of every native endpoint-builder input that matters to the UI
/// protocol.  The production binder is deliberately a small conversion from this plan to iroh's
/// 1.0.3 builder API.
#[derive(Clone)]
pub(crate) struct IrohEndpointBuildPlan {
    route: IrohEndpointRoute,
    relay_route: IrohRelayRoute,
    bind_addr: Option<SocketAddr>,
    keep_alive_interval: Duration,
    max_idle_timeout: Duration,
    max_concurrent_connections: usize,
    max_concurrent_bidirectional_streams: u32,
    max_concurrent_unidirectional_streams: u32,
    stream_accept_timeout: Duration,
    max_tls_tickets: usize,
}

impl IrohEndpointBuildPlan {
    pub(crate) fn public(relays: IrohRelayConfig) -> Self {
        let relay_route = match relays {
            IrohRelayConfig::N0Preset => IrohRelayRoute::N0,
            IrohRelayConfig::Explicit(relays) => IrohRelayRoute::Explicit(relays),
        };
        Self::new(IrohEndpointRoute::RelayOnly, relay_route, None)
    }

    #[cfg(test)]
    pub(crate) fn loopback(bind_addr: SocketAddr) -> Self {
        Self::new(
            IrohEndpointRoute::DirectLoopback,
            IrohRelayRoute::Disabled,
            Some(bind_addr),
        )
    }

    fn new(
        route: IrohEndpointRoute,
        relay_route: IrohRelayRoute,
        bind_addr: Option<SocketAddr>,
    ) -> Self {
        Self {
            route,
            relay_route,
            bind_addr,
            keep_alive_interval: KEEP_ALIVE_INTERVAL,
            max_idle_timeout: MAX_IDLE_TIMEOUT,
            max_concurrent_connections: MAX_CONCURRENT_CONNECTIONS,
            max_concurrent_bidirectional_streams: MAX_CONCURRENT_BIDIRECTIONAL_STREAMS,
            max_concurrent_unidirectional_streams: MAX_CONCURRENT_UNIDIRECTIONAL_STREAMS,
            stream_accept_timeout: STREAM_ACCEPT_TIMEOUT,
            // iroh documents this cache as the facility used for 0-RTT connections.  Keeping it
            // empty ensures this replay-sensitive protocol never obtains a resumable ticket.
            max_tls_tickets: 0,
        }
    }

    #[cfg(test)]
    pub(crate) fn alpn(&self) -> &[u8] {
        UI_IROH_ALPN
    }

    #[cfg(test)]
    pub(crate) fn route(&self) -> &IrohEndpointRoute {
        &self.route
    }

    #[cfg(test)]
    pub(crate) fn relay_urls(&self) -> Vec<String> {
        match &self.relay_route {
            IrohRelayRoute::Explicit(relays) => relays.iter().map(ToString::to_string).collect(),
            IrohRelayRoute::N0 | IrohRelayRoute::Disabled => Vec::new(),
        }
    }

    #[cfg(test)]
    pub(crate) fn bind_addr(&self) -> Option<SocketAddr> {
        self.bind_addr
    }

    #[cfg(test)]
    pub(crate) fn keep_alive_interval(&self) -> Duration {
        self.keep_alive_interval
    }

    #[cfg(test)]
    pub(crate) fn max_idle_timeout(&self) -> Duration {
        self.max_idle_timeout
    }

    #[cfg(test)]
    pub(crate) fn max_concurrent_connections(&self) -> usize {
        self.max_concurrent_connections
    }

    #[cfg(test)]
    pub(crate) fn max_concurrent_bidirectional_streams(&self) -> u32 {
        self.max_concurrent_bidirectional_streams
    }

    #[cfg(test)]
    pub(crate) fn max_concurrent_unidirectional_streams(&self) -> u32 {
        self.max_concurrent_unidirectional_streams
    }

    #[cfg(test)]
    pub(crate) fn stream_accept_timeout(&self) -> Duration {
        self.stream_accept_timeout
    }

    #[cfg(test)]
    pub(crate) fn max_tls_tickets(&self) -> usize {
        self.max_tls_tickets
    }
}

#[derive(Clone, Debug)]
pub(crate) struct IrohEndpointStatus {
    pub(crate) endpoint_id: String,
    pub(crate) endpoint_ticket: Option<String>,
    pub(crate) relay_ready: bool,
}

pub(crate) fn initial_endpoint_status(
    plan: &IrohEndpointBuildPlan,
    endpoint_id: String,
    endpoint_ticket: String,
) -> IrohEndpointStatus {
    let endpoint_ticket = match plan.route {
        IrohEndpointRoute::RelayOnly => None,
        IrohEndpointRoute::DirectLoopback => Some(endpoint_ticket),
    };
    IrohEndpointStatus {
        endpoint_id,
        endpoint_ticket,
        relay_ready: false,
    }
}

pub(crate) fn publish_online_status(status: &mut IrohEndpointStatus, endpoint_ticket: String) {
    status.endpoint_ticket = Some(endpoint_ticket);
    status.relay_ready = true;
}

pub(crate) struct RunningIrohEndpoint {
    pub(crate) status: watch::Receiver<IrohEndpointStatus>,
    pub(crate) task: JoinHandle<AgentResult<()>>,
}

#[derive(Clone)]
pub(crate) struct IrohEndpointFactory {
    identity_store: IrohIdentityStore,
    plan: IrohEndpointBuildPlan,
    hub: Arc<UiHub>,
    pairing_service: Arc<UiPairingService>,
}

impl IrohEndpointFactory {
    pub(crate) fn public(
        identity_store: IrohIdentityStore,
        relays: IrohRelayConfig,
        hub: Arc<UiHub>,
        pairing_service: Arc<UiPairingService>,
    ) -> Self {
        Self::new(
            identity_store,
            IrohEndpointBuildPlan::public(relays),
            hub,
            pairing_service,
        )
    }

    pub(crate) fn new(
        identity_store: IrohIdentityStore,
        plan: IrohEndpointBuildPlan,
        hub: Arc<UiHub>,
        pairing_service: Arc<UiPairingService>,
    ) -> Self {
        Self {
            identity_store,
            plan,
            hub,
            pairing_service,
        }
    }

    pub(crate) async fn start(
        &self,
        cancellation: CancellationToken,
    ) -> AgentResult<RunningIrohEndpoint> {
        let identity = self.identity_store.load_or_create().map_err(|_| {
            AgentError::new(
                ErrorCode::PairingStorageFailed,
                "Iroh endpoint identity is unavailable",
            )
        })?;
        let endpoint = bind_endpoint(identity.secret_key().clone(), &self.plan).await?;
        let status = initial_endpoint_status(
            &self.plan,
            endpoint.id().to_string(),
            EndpointTicket::new(endpoint.addr()).to_string(),
        );
        let (status_tx, status_rx) = watch::channel(status);
        let task = tokio::spawn(run_endpoint(
            endpoint,
            self.plan.clone(),
            self.hub.clone(),
            self.pairing_service.clone(),
            cancellation,
            status_tx,
        ));
        Ok(RunningIrohEndpoint {
            status: status_rx,
            task,
        })
    }
}

async fn bind_endpoint(
    secret_key: iroh::SecretKey,
    plan: &IrohEndpointBuildPlan,
) -> AgentResult<Endpoint> {
    let idle_timeout = plan
        .max_idle_timeout
        .try_into()
        .map_err(|_| endpoint_error())?;
    let transport_config = QuicTransportConfig::builder()
        .keep_alive_interval(plan.keep_alive_interval)
        .max_idle_timeout(Some(idle_timeout))
        .max_concurrent_bidi_streams(plan.max_concurrent_bidirectional_streams.into())
        .max_concurrent_uni_streams(plan.max_concurrent_unidirectional_streams.into())
        .build();
    let builder = Endpoint::builder(presets::Minimal)
        .secret_key(secret_key)
        .alpns(vec![UI_IROH_ALPN.to_vec()])
        .clear_address_lookup()
        .max_tls_tickets(plan.max_tls_tickets)
        .transport_config(transport_config)
        .clear_ip_transports();
    let builder = match plan.route {
        IrohEndpointRoute::RelayOnly => builder.proxy_from_env(),
        IrohEndpointRoute::DirectLoopback => builder,
    };
    let builder = match &plan.relay_route {
        IrohRelayRoute::N0 => builder.relay_mode(RelayMode::Default),
        IrohRelayRoute::Explicit(relays) => builder.relay_mode(RelayMode::custom(relays.clone())),
        IrohRelayRoute::Disabled => builder.relay_mode(RelayMode::Disabled),
    };
    let builder = match plan.bind_addr {
        Some(bind_addr) => builder.bind_addr(bind_addr).map_err(|_| endpoint_error())?,
        None => builder,
    };
    builder.bind().await.map_err(|_| endpoint_error())
}

async fn run_endpoint(
    endpoint: Endpoint,
    plan: IrohEndpointBuildPlan,
    hub: Arc<UiHub>,
    pairing_service: Arc<UiPairingService>,
    cancellation: CancellationToken,
    status: watch::Sender<IrohEndpointStatus>,
) -> AgentResult<()> {
    let relay_status_endpoint = endpoint.clone();
    let relay_status_cancellation = cancellation.clone();
    let relay_status_tx = status.clone();
    let waits_for_relay = plan.route == IrohEndpointRoute::RelayOnly;
    let relay_status = AbortOnDropHandle::new(tokio::spawn(async move {
        if !waits_for_relay {
            return;
        }
        let mut watcher = relay_status_endpoint.home_relay_status();
        loop {
            let connected = watcher.get().iter().any(|status| status.is_connected());
            publish_relay_status(&relay_status_tx, &relay_status_endpoint, connected);
            tokio::select! {
                _ = relay_status_cancellation.cancelled() => return,
                next = watcher.updated() => match next {
                    Ok(_) => {}
                    Err(_) => return,
                },
            }
        }
    }));
    let permits = Arc::new(Semaphore::new(plan.max_concurrent_connections));
    let mut clients = JoinSet::new();

    loop {
        tokio::select! {
            _ = cancellation.cancelled() => break,
            incoming = endpoint.accept() => {
                let Some(incoming) = incoming else {
                    break;
                };
                let Ok(permit) = permits.clone().try_acquire_owned() else {
                    incoming.refuse();
                    continue;
                };
                let hub = hub.clone();
                let pairing_service = pairing_service.clone();
                let plan = plan.clone();
                let connection_cancellation = cancellation.clone();
                clients.spawn(async move {
                    let _permit = permit;
                    if let Err(error) = handle_connection(
                        incoming,
                        plan,
                        hub,
                        pairing_service,
                        connection_cancellation,
                    )
                    .await
                    {
                        tracing::debug!(error = %error, "iroh UI connection ended");
                    }
                });
            }
            completed = clients.join_next(), if !clients.is_empty() => {
                if completed.is_some_and(|result| result.is_err()) {
                    tracing::debug!("iroh UI connection task failed");
                }
            }
        }
    }

    cancellation.cancel();
    let _ = timeout(ENDPOINT_SHUTDOWN_TIMEOUT, async {
        while clients.join_next().await.is_some() {}
    })
    .await;
    clients.shutdown().await;
    drop(relay_status);
    let _ = timeout(ENDPOINT_SHUTDOWN_TIMEOUT, endpoint.close()).await;
    Ok(())
}

fn publish_relay_status(
    status: &watch::Sender<IrohEndpointStatus>,
    endpoint: &Endpoint,
    connected: bool,
) {
    let ticket = connected.then(|| EndpointTicket::new(endpoint.addr()).to_string());
    status.send_if_modified(|current| {
        if connected {
            let ticket = ticket.clone().expect("connected relay has a ticket");
            if current.relay_ready && current.endpoint_ticket.as_deref() == Some(ticket.as_str()) {
                return false;
            }
            publish_online_status(current, ticket);
            true
        } else if current.relay_ready || current.endpoint_ticket.is_some() {
            current.relay_ready = false;
            current.endpoint_ticket = None;
            true
        } else {
            false
        }
    });
}

async fn handle_connection(
    incoming: iroh::endpoint::Incoming,
    plan: IrohEndpointBuildPlan,
    hub: Arc<UiHub>,
    pairing_service: Arc<UiPairingService>,
    endpoint_cancellation: CancellationToken,
) -> AgentResult<()> {
    // Awaiting `Incoming` completes the cryptographic handshake.  We intentionally never call
    // `into_0rtt`, so pairing and agent operations cannot be replayed.
    let connection = tokio::select! {
        _ = endpoint_cancellation.cancelled() => return Ok(()),
        connection = incoming => connection.map_err(|_| endpoint_error())?,
    };
    let remote_endpoint = connection.remote_id().to_string();
    let peer = UiChannelPeer {
        source: SourceIdentity::IrohEndpoint(remote_endpoint.clone()),
        binding: ClientBinding::IrohEndpoint {
            endpoint_id: remote_endpoint,
        },
    };
    let (send, recv) = match timeout(plan.stream_accept_timeout, connection.accept_bi()).await {
        Ok(Ok(streams)) => streams,
        Ok(Err(_)) | Err(_) => {
            connection.close(CLOSE_STREAM_TIMEOUT.into(), b"stream_timeout");
            return Ok(());
        }
    };
    let connection_for_extra_streams = connection.clone();
    let closed = CancellationToken::new();
    let closed_on_endpoint_shutdown = closed.clone();
    let endpoint_shutdown = AbortOnDropHandle::new(tokio::spawn(async move {
        endpoint_cancellation.cancelled().await;
        closed_on_endpoint_shutdown.cancel();
    }));
    let extra_streams_closed = closed.clone();
    let extra_streams = AbortOnDropHandle::new(tokio::spawn(async move {
        tokio::select! {
            _ = extra_streams_closed.cancelled() => {}
            extra = connection_for_extra_streams.accept_bi() => {
                if extra.is_ok() {
                    connection_for_extra_streams.close(CLOSE_STREAM_LIMIT.into(), b"stream_limit");
                }
            }
        }
    }));
    let (incoming_tx, incoming_rx) = mpsc::channel(UI_CLIENT_QUEUE_CAPACITY);
    let (outgoing_tx, outgoing_rx) = mpsc::channel(UI_CLIENT_QUEUE_CAPACITY);
    let adapter = AbortOnDropHandle::new(tokio::spawn(adapt_iroh_stream(
        send,
        recv,
        incoming_tx,
        outgoing_rx,
        closed.clone(),
    )));
    let client_id = hub.next_client_id();
    let mut channel = UiChannel {
        incoming: incoming_rx,
        outgoing: outgoing_tx,
        closed: closed.clone(),
    };
    let approved_client = authenticate_authenticated_channel(
        &mut channel,
        pairing_service,
        ConnectionId::new(client_id),
        &peer,
    )
    .await?;
    if let Some(approved_client) = approved_client {
        hub.attach_authenticated_with_id(client_id, peer, approved_client, channel)
            .await?;
    }
    closed.cancel();
    drop(endpoint_shutdown);
    drop(extra_streams);
    let adapter_result = adapter.await.map_err(|_| endpoint_error())?;
    if adapter_result.is_err() {
        connection.close(CLOSE_PROTOCOL.into(), b"invalid_request");
    } else {
        connection.close(0_u32.into(), b"closed");
    }
    Ok(())
}

async fn adapt_iroh_stream(
    mut send: iroh::endpoint::SendStream,
    mut recv: iroh::endpoint::RecvStream,
    incoming: mpsc::Sender<String>,
    mut outgoing: mpsc::Receiver<String>,
    closed: CancellationToken,
) -> AgentResult<()> {
    let mut decoder = JsonFrameDecoder::new(FrameLimit::Handshake);
    let mut authenticated = false;
    let mut waiting_for_handshake_response = false;
    let mut buffer = Vec::new();
    let mut buffer_offset = 0_usize;
    let mut read_buffer = [0_u8; 8 * 1024];
    let mut terminal_unauthenticated_response_written = false;

    let result = loop {
        if !waiting_for_handshake_response && buffer_offset < buffer.len() {
            let chunk = &buffer[buffer_offset..];
            match decoder.push_one(chunk).map_err(|_| protocol_error())? {
                Some((record, suffix)) => {
                    buffer_offset += chunk.len() - suffix.len();
                    serde_json::from_str::<Value>(&record).map_err(|_| protocol_error())?;
                    if incoming.send(record).await.is_err() {
                        break Ok(());
                    }
                    if !authenticated {
                        waiting_for_handshake_response = true;
                    }
                    continue;
                }
                None => {
                    buffer.clear();
                    buffer_offset = 0;
                    continue;
                }
            }
        }
        if buffer_offset == buffer.len() {
            buffer.clear();
            buffer_offset = 0;
        }

        tokio::select! {
            _ = closed.cancelled() => {
                if let Ok(text) = outgoing.try_recv() {
                    write_framed_json(&mut send, &text, authenticated).await?;
                    terminal_unauthenticated_response_written = !authenticated
                        && !is_successful_handshake_response(&text);
                }
                break Ok(());
            }
            outgoing = outgoing.recv() => {
                let Some(text) = outgoing else {
                    break Ok(());
                };
                write_framed_json(&mut send, &text, authenticated).await?;
                if !authenticated {
                    if is_successful_handshake_response(&text) {
                        authenticated = true;
                        waiting_for_handshake_response = false;
                        decoder.set_limit(FrameLimit::Authenticated);
                    } else {
                        terminal_unauthenticated_response_written = true;
                    }
                }
            }
            read = recv.read(&mut read_buffer), if !waiting_for_handshake_response => {
                let count = read.map_err(|_| protocol_error())?;
                let Some(count) = count else {
                    decoder.finish().map_err(|_| protocol_error())?;
                    break Ok(());
                };
                buffer.extend_from_slice(&read_buffer[..count]);
            }
        }
    };
    closed.cancel();
    let _ = send.shutdown().await;
    if terminal_unauthenticated_response_written {
        // `shutdown` only starts the QUIC FIN.  Keep the connection alive long
        // enough for the browser to acknowledge the final rejection frame;
        // otherwise the connection close below can stop retransmission first.
        let _ = timeout(TERMINAL_HANDSHAKE_DRAIN_TIMEOUT, send.stopped()).await;
    }
    result
}

async fn write_framed_json(
    send: &mut iroh::endpoint::SendStream,
    text: &str,
    authenticated: bool,
) -> AgentResult<()> {
    let limit = if authenticated {
        FrameLimit::Authenticated
    } else {
        FrameLimit::Handshake
    };
    let frame = encode_json_frame(text, limit).map_err(|_| protocol_error())?;
    send.write_all(&frame).await.map_err(|_| protocol_error())
}

fn is_successful_handshake_response(text: &str) -> bool {
    serde_json::from_str::<Value>(text)
        .ok()
        .and_then(|value| value.get("kind").and_then(Value::as_str).map(str::to_owned))
        .is_some_and(|kind| kind == "client.paired" || kind == "client.authenticated")
}

fn endpoint_error() -> AgentError {
    AgentError::new(ErrorCode::InvalidMessage, "Iroh endpoint is unavailable")
}

fn protocol_error() -> AgentError {
    AgentError::new(ErrorCode::InvalidMessage, "Iroh UI stream protocol error")
}