package main

import (
	"net/http"
	"net/rpc"
	"net/rpc/jsonrpc"

	"github.com/gorilla/websocket"
)

type Args struct {
	A int
	B int
}

type Arith int

func (t *Arith) Multiply(args *Args, reply *int) error {
	*reply = args.A * args.B
	return nil
}

type WebSocketConn struct {
	conn *websocket.Conn
}

func (wsc *WebSocketConn) Read(p []byte) (n int, err error) {
	_, msg, err := wsc.conn.ReadMessage()
	if err != nil {
		return 0, err
	}
	copy(p, msg)
	return len(msg), nil
}

func (wsc *WebSocketConn) Write(p []byte) (n int, err error) {
	err = wsc.conn.WriteMessage(websocket.TextMessage, p)
	if err != nil {
		return 0, err
	}
	return len(p), nil
}

func (wsc *WebSocketConn) Close() error {
	return wsc.conn.Close()
}

var (
	upgrader = websocket.Upgrader{}
	rpcServer = rpc.NewServer()
	
)

// WebSocket 处理函数
func wsHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		http.Error(w, "Failed to upgrade connection", http.StatusBadRequest)
		return
	}

	go rpcServer.ServeCodec(jsonrpc.NewServerCodec(&WebSocketConn{conn: conn}))
}

func main() {
	rpcServer.Register(new(Arith))
	http.HandleFunc("/ws", wsHandler)
	http.ListenAndServe(":7000", nil)
}