// Copyright 2026 Leo Cheng
// SPDX-License-Identifier: Apache-2.0

///|
/// A QUIC connection driven by one task: the client half of what `moonquic` states and this
/// package puts on a datagram socket. Reading is a poll, because a connection that also
/// carries control and QPACK streams cannot hand every arrival to one caller.
pub struct Quic {
  socket : Socket
  to : @socket.Addr
  client : @conn.Client
  mut bi : UInt64
  mut uni : UInt64
  mut dead : Bool
  // Streams the peer has sent something on. The codec reports an unknown stream as ended,
  // because it has no record of one, so a stream this side only wrote on is not ended yet.
  seen : Map[UInt64, Bool]
  // Octets handed to `write` on each stream. The send loop keeps its own record of a stream it
  // has carried data for; a stream opened and never written on has none, and still has to be
  // reset with a final size rather than not at all.
  sent : Map[UInt64, UInt64]
  // How the peer ended each stream abruptly: the code of a reset it sent, or of the stop it
  // asked for. A stream the peer ended with a FIN is not in here.
  errors : Map[UInt64, QuicRefused]
  // Streams this side has abandoned as a sender, with the code it sent.
  abandoned : Map[UInt64, UInt64]
}

///|
/// The clock a QUIC connection counts in: microseconds since the epoch, which is the unit the
/// recovery, idle and closing-period timers are measured in (RFC 9002 §6). Reading it in
/// milliseconds leaves every timer a thousand times too long, so a lost packet is never sent
/// again.
pub fn now() -> Int64 {
  (@env.now() * 1000).reinterpret_as_int64()
}

///|
/// The clock a certificate is judged at: seconds since the epoch.
fn wall() -> Int64 {
  (@env.now() / 1000).reinterpret_as_int64()
}

///|
fn entropy(n : Int32) -> Bytes {
  guard @env.rand(n) is Some(bytes) else { abort("no system entropy") }
  bytes
}

///|
/// How long a QUIC dial gives one address before it tries the next one as well: the
/// connection attempt delay of RFC 8305 §5, at the 250 ms §8 recommends.
pub let attempt_delay : Int32 = 250

///|
/// Connect to `host:port` and shake hands. The connection is returned once the handshake is
/// confirmed; `tls` offers `h3` or whatever ALPN the peer is expected to choose.
///
/// A name with both an IPv6 and an IPv4 address is dialled as RFC 8305 says: IPv6 first, IPv4
/// as well once `attempt_delay` milliseconds pass without an answer or the first attempt fails,
/// and the first handshake to complete is kept. A path that drops datagrams costs one delay,
/// not a handshake timeout. A name with one address is dialled as it always was.
pub async fn dial(
  host : String,
  port~ : Int32,
  tls~ : @tls.Client,
  idle~ : UInt64 = 30000UL,
  attempt_delay~ : Int32 = attempt_delay,
) -> Quic raise QuicRefused {
  race(addresses(host, port), host, tls, idle, attempt_delay)
}

///|
/// `dial` to addresses already looked up, such as the ones a service registry gives: they are
/// tried in the order given, with the same delay between attempts, and `host` is the name the
/// certificate is checked against.
pub async fn dial_addrs(
  addrs : ArrayView[@socket.Addr],
  host~ : String,
  tls~ : @tls.Client,
  idle~ : UInt64 = 30000UL,
  attempt_delay~ : Int32 = attempt_delay,
) -> Quic raise QuicRefused {
  guard !addrs.is_empty() else {
    raise QuicRefused::Broken("no address to dial")
  }
  race(addrs.to_owned(), host, tls, idle, attempt_delay)
}

///|
/// The addresses `host` names, IPv6 first (RFC 8305 §4), one of each family at most.
async fn addresses(
  host : String,
  port : Int32,
) -> Array[@socket.Addr] raise QuicRefused {
  let found : Array[@socket.Addr] = []
  for family in [@socket.OnlyV6, @socket.OnlyV4] {
    let addr = @socket.Addr::resolve(host, port~, protocol=family) catch {
      _ => continue
    }
    if !found.contains(addr) {
      found.push(addr)
    }
  }
  if found.is_empty() {
    // Neither family answered on its own; the plain lookup says why.
    found.push(
      @socket.Addr::resolve(host, port~) catch {
        e => raise QuicRefused::Broken("\{e}")
      },
    )
  }
  found
}

///|
/// Shake hands with the first of `addrs` to answer. The first attempt starts at once; each
/// further one starts when `delay` has passed since the one before, or sooner when an attempt
/// fails. The first handshake to complete wins and the others are closed; when every attempt
/// fails, the last failure is raised.
async fn race(
  addrs : Array[@socket.Addr],
  host : String,
  tls : @tls.Client,
  idle : UInt64,
  delay : Int32,
) -> Quic raise QuicRefused {
  if addrs.length() == 1 {
    return attempt(addrs[0], host, tls, idle)
  }
  let failures : Array[QuicRefused] = []
  let failed = @cond_var.Cond()
  @async.with_task_group(g => {
    for i, to in addrs {
      g.spawn_bg(allow_failure=true, () => {
        let link = attempt(to, host, tls, idle) catch {
          e => {
            failures.push(e)
            failed.broadcast()
            return
          }
        }
        g.return_immediately(link)
      })
      if i + 1 < addrs.length() {
        let before = failures.length()
        ignore(
          @async.with_timeout_opt(delay, () => {
            while failures.length() == before {
              failed.wait()
            }
          }),
        )
      }
    }
    while failures.length() < addrs.length() {
      failed.wait()
    }
    raise failures[failures.length() - 1]
  }) catch {
    e => raise refusal(e)
  }
}

///|
/// An error out of a task group, as the refusal it carries.
fn refusal(e : Error) -> QuicRefused {
  match e {
    QuicRefused::Timeout => Timeout
    Broken(why) => Broken(why)
    Closed => Closed
    Reset(code) => Reset(code)
    Stopped(code) => Stopped(code)
    _ => Broken("\{e}")
  }
}

///|
/// One handshake with one address. The socket is closed on the way out of a handshake that
/// did not complete, including one cancelled because another address answered first.
async fn attempt(
  to : @socket.Addr,
  host : String,
  tls : @tls.Client,
  idle : UInt64,
) -> Quic raise QuicRefused {
  // The socket has to be of the same family as the peer, or the kernel refuses the datagram.
  let socket = bind(host=if to.is_ipv6() { "::" } else { "0.0.0.0" }, port=0) catch {
    e => raise QuicRefused::Broken("\{e}")
  }
  // A handshake that does not complete, or that another address won, leaves no socket behind.
  errdefer socket.close()
  let scid = entropy(8)
  let dcid = entropy(8)
  let params = @conn.Limits::new(
    idle_timeout=idle,
    data=8_000_000UL,
    stream_data_bidi_local=4_000_000UL,
    stream_data_bidi_remote=4_000_000UL,
    stream_data_uni=1_000_000UL,
    streams_bidi=128UL,
    streams_uni=16UL,
  ).params()
  params.push({ kind: @conn.Kind::InitialScid, value: scid, })
  let client = @conn.Client::new(
    tls,
    host,
    @conn.encode(params[:]),
    entropy,
    dcid,
    scid,
    wall(),
  ) catch {
    e => raise QuicRefused::Broken("\{e}")
  }
  let link : Quic = {
    socket,
    to,
    client,
    bi: 0UL,
    uni: 2UL,
    dead: false,
    seen: Map([]),
    sent: Map([]),
    errors: Map([]),
    abandoned: Map([]),
  }
  let mut rounds = 0
  while !link.client.is_connected() {
    if rounds > 600 {
      raise QuicRefused::Timeout
    }
    let _ = link.step(timeout=20)
    rounds += 1
  }
  link
}

///|
/// What a QUIC connection will not do.
#derive(Eq, Debug)
pub(all) suberror QuicRefused {
  /// The handshake did not complete in the time this waited.
  Timeout
  /// The socket or the connection under it failed.
  Broken(String)
  /// The connection is over.
  Closed
  /// The stream was abandoned with this application code: the peer sent RESET_STREAM when
  /// this side was reading it, or this side sent one itself (RFC 9000 §19.4).
  Reset(UInt64)
  /// The peer sent STOP_SENDING with this code, which says it will read no more of the
  /// stream; sending on it is over (RFC 9000 §19.5).
  Stopped(UInt64)
}

///|
pub extend QuicRefused with Eq::{equal, not_equal}

///|
pub extend QuicRefused with Debug::{to_repr}

///|
/// Send what the connection has queued and take one datagram in, which is one turn of a
/// handshake or of an exchange. False when nothing arrived within `timeout`.
pub async fn Quic::step(
  self : Quic,
  timeout~ : Int32 = 20,
) -> Bool raise QuicRefused {
  if self.dead {
    return false
  }
  for datagram in self.client.poll_out(now()) {
    self.socket.send(self.to, datagram) catch {
      e => raise QuicRefused::Broken("\{e}")
    }
  }
  let arrived = self.socket.recv(timeout~) catch {
    e => raise QuicRefused::Broken("\{e}")
  }
  guard arrived is Some((bytes, _)) else { return false }
  self.client.recv(bytes, now()) catch {
    e => {
      self.dead = true
      raise QuicRefused::Broken("\{e}")
    }
  }
  for event in self.client.take_events() {
    match event {
      @conn.Event::Readable(id) => self.seen[id] = true
      @conn.Event::Reset(id, code) => self.errors[id] = QuicRefused::Reset(code)
      // §3.5: a STOP_SENDING leaves the sender no choice but to abandon the stream, with the
      // code the peer gave it.
      @conn.Event::Stopped(id, code) => {
        self.errors[id] = QuicRefused::Stopped(code)
        ignore(self.abandon(id, code))
      }
      @conn.Event::Closed(_) => self.dead = true
      _ => ()
    }
  }
  true
}

///|
/// Put `bytes` on `stream`, and end it when `fin`. Raises `Stopped` when the peer has asked
/// this side to stop sending on the stream, and `Reset` when this side abandoned it.
pub fn Quic::write(
  self : Quic,
  stream : UInt64,
  bytes : Bytes,
  fin~ : Bool = false,
) -> Unit raise QuicRefused {
  if self.dead {
    raise QuicRefused::Closed
  }
  if self.errors.get(stream) is Some(QuicRefused::Stopped(code)) {
    raise QuicRefused::Stopped(code)
  }
  if self.abandoned.get(stream) is Some(code) {
    raise QuicRefused::Reset(code)
  }
  self.client.write(stream, bytes, fin~)
  if bytes.length() > 0 {
    let at = self.sent.get(stream).unwrap_or(0UL)
    self.sent[stream] = at + bytes.length().to_uint64()
  }
}

///|
/// The octets that arrived on `stream` since the last call; empty when none have. What arrived
/// before a reset is still handed over first: RFC 9000 §3.2 keeps the octets received up to the
/// final size readable, and it is the end that is an error. An empty read raises `Reset` once
/// the peer has abandoned the stream.
pub fn Quic::take(self : Quic, stream : UInt64) -> Bytes raise QuicRefused {
  let data = self.client.read(stream)
  if data.length() > 0 {
    self.seen[stream] = true
    return data
  }
  if self.errors.get(stream) is Some(QuicRefused::Reset(code)) {
    raise QuicRefused::Reset(code)
  }
  data
}

///|
/// How the peer ended `stream` abruptly, if it did: `Reset` for a stream this side was reading,
/// `Stopped` for one it was writing. `None` while the stream is live, and after a clean FIN —
/// which is what `ended` reports.
pub fn Quic::refused(self : Quic, stream : UInt64) -> QuicRefused? {
  self.errors.get(stream)
}

///|
/// Whether the peer has ended `stream`, with every octet of it read. A stream the peer reset is
/// over too, however much of it arrived.
pub fn Quic::ended(self : Quic, stream : UInt64) -> Bool {
  if self.errors.get(stream) is Some(QuicRefused::Reset(_)) {
    return true
  }
  self.seen.get(stream).unwrap_or(false) && self.client.finished(stream)
}

///|
/// Abandon `stream` as a sender: send RESET_STREAM carrying `code`, and drop what the send loop
/// still holds for it (RFC 9000 §3.5, §19.4). Zero, the default, is NO_ERROR (§20.1). Writing
/// on the stream afterwards raises `Reset`. A stream this side only reads may be reset too: it
/// closes this side's sending half, which is what a reply abandoned half way needs.
pub fn Quic::reset(
  self : Quic,
  stream : UInt64,
  code~ : UInt64 = 0UL,
) -> Unit raise QuicRefused {
  if self.dead {
    raise QuicRefused::Closed
  }
  if self.abandoned.contains(stream) {
    return
  }
  if !self.abandon(stream, code) {
    raise QuicRefused::Closed
  }
  self.abandoned[stream] = code
}

///|
/// Ask the peer to stop sending on `stream`, with STOP_SENDING carrying `code` (RFC 9000
/// §19.5): the octets already on their way are dropped when they arrive, and `take` raises
/// `Reset` if the peer answers with one. Zero, the default, is NO_ERROR (§20.1).
pub fn Quic::stop(
  self : Quic,
  stream : UInt64,
  code~ : UInt64 = 0UL,
) -> Unit raise QuicRefused {
  if self.dead {
    raise QuicRefused::Closed
  }
  let frame : @frame.Frame = @frame.Frame::StopSending(
    id=stream,
    error_code=code,
  )
  if !self.client.send_on(0UL, [frame], now()) {
    raise QuicRefused::Closed
  }
}

///|
/// Queue a RESET_STREAM for `stream` and cancel what is queued on it, answering whether the
/// connection could take the frame. The send loop queues the frame itself while it still holds
/// the stream; a stream whose data has all been acknowledged has been released there, and then
/// the frame is queued here with the final size this side wrote.
fn Quic::abandon(self : Quic, stream : UInt64, code : UInt64) -> Bool {
  if self.client.buffered(stream) > 0UL {
    self.client.sender.abandon(stream, code)
    return true
  }
  let end = self.sent.get(stream).unwrap_or(0UL)
  let frame : @frame.Frame = @frame.Frame::ResetStream(
    id=stream,
    error_code=code,
    final_size=end,
  )
  self.client.send_on(0UL, [frame], now())
}

///|
/// The next bidirectional stream this side opens: 0, 4, 8, … as RFC 9000 §2.1 numbers them.
pub fn Quic::open_bi(self : Quic) -> UInt64 {
  let id = self.bi
  self.bi += 4UL
  id
}

///|
/// The next unidirectional stream, from 2 upwards.
pub fn Quic::open_uni(self : Quic) -> UInt64 {
  let id = self.uni
  self.uni += 4UL
  id
}

///|
/// The ALPN the handshake settled on.
pub fn Quic::alpn(self : Quic) -> String? {
  self.client.alpn()
}

///|
/// Close the socket the connection runs on.
pub fn Quic::close(self : Quic) -> Unit {
  self.dead = true
  self.socket.close()
}