package golem:websocket@1.5.0;

interface client {
  use wasi:io/poll@0.2.3.{pollable};

  variant error {
    connection-failure(string),
    send-failure(string),
    receive-failure(string),
    protocol-error(string),
    closed(option<close-info>),
    other(string),
  }

  record close-info {
    code: u16,
    reason: string,
  }

  /// A WebSocket message — text or binary
  variant message {
    text(string),
    binary(list<u8>),
  }

  /// A WebSocket connection resource
  resource websocket-connection {
    /// Connect to a WebSocket server at the given URL (ws:// or wss://)
    /// Optional headers for auth, subprotocols, etc.
    connect: static func(
      url: string,
      headers: option<list<tuple<string, string>>>
    ) -> result<websocket-connection, error>;

    /// Send a message (text or binary)
    send: func(message: message) -> result<_, error>;

    /// Receive the next message (blocks until available)
    receive: func() -> result<message, error>;

    /// Receive the next message with a timeout in milliseconds.
    /// Returns none if the timeout expires before a message arrives.
    receive-with-timeout: func(timeout-ms: u64) -> result<option<message>, error>;

    /// Send a close frame with optional code and reason
    close: func(code: option<u16>, reason: option<string>) -> result<_, error>;

    /// Returns a pollable that resolves when a message is available to read
    subscribe: func() -> pollable;
  }
}

@agent()
class ExampleAgent extends BaseAgent {
  async run(): Promise<void> {
    return new Promise((resolve, reject) => {
      const ws = new WebSocket("wss://example.com/chat");
  
      ws.onopen = () => {
        console.log("Connected");
        ws.send("Hello, server!");
      };
  
      ws.onmessage = (event: MessageEvent) => {
        if (typeof event.data === "string") {
          console.log("Text:", event.data);
        } else {
          console.log("Binary:", new Uint8Array(event.data));
        }
      };
  
      ws.onerror = () => reject(new Error("WebSocket error"));
      ws.onclose = (event: CloseEvent) => {
        console.log(`Closed [${event.code}] "${event.reason}"`);
        resolve();
      };
    });
  }
}

#[agent_implementation]
impl ExampleAgent for ExampleAgentImpl {
    async fn run() -> Result<(), WebSocketError> {
        let ws = WebsocketConnection::connect("wss://example.com/chat", None)?;
        println!("Connected");
    
        ws.send(&WebSocketMessage::Text("Hello, server!".to_string()))?;
    
        loop {
            match ws.receive().await {
                Ok(WebSocketMessage::Text(text)) => println!("Text: {text}"),
                Ok(WebSocketMessage::Binary(data)) => println!("Binary: {data:?}"),
                Err(WebSocketError::Closed(info)) => {
                    if let Some(info) = info {
                        println!("Closed [{}] \"{}\"", info.code, info.reason);
                    }
                    break;
                }
                Err(err) => return Err(err),
            }
        }
    
        Ok(())
    }
}

case class ExampleAgentImpl() extends ExampleAgent {
  def run(): Future[Unit] = {
    val done = Promise[Unit]()
    val ws = new WebSocket("wss://example.com/chat")
  
    ws.onopen = { (_: Event) =>
      println("Connected")
      ws.send("Hello, server!")
    }
  
    ws.onmessage = { (event: MessageEvent) =>
      event.data match {
        case text: String => println(s"Text: $text")
        case other        => println(s"Binary: $other")
      }
    }
  
    ws.onerror = { (_: Event) =>
      done.tryFailure(new Exception("WebSocket error"))
    }
  
    ws.onclose = { (event: CloseEvent) =>
      println(s"Closed [${event.code}] \"${event.reason}\"")
      done.trySuccess(())
    }
  
    done.future
  }
}

pub fn ExampleAgent::run(self : Self) -> Unit raise @common.AgentError {
  let conn = match @websocket_client.WebsocketConnection::connect(
    "wss://example.com/chat", None,
  ) {
    Ok(c) => c
    Err(e) => raise @common.AgentError::InvalidInput("Connect failed: \{e}")
  }
  println("Connected")
  match conn.send(Text("Hello, server!")) {
    Ok(_) => ()
    Err(e) => raise @common.AgentError::InvalidInput("Send failed: \{e}")
  }
  while true {
    match conn.receive() {
      Ok(Text(msg)) => println("Text: \{msg}")
      Ok(Binary(data)) => println("Binary: \{data.length()} bytes")
      Err(Closed(Some(info))) => {
        println("Closed [\{info.code}] \"\{info.reason}\"")
        break
      }
      Err(Closed(None)) => break
      Err(e) => raise @common.AgentError::InvalidInput("Receive failed: \{e}")
    }
  }
  conn.drop()
}