package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"agent/src/message"
	"agent/src/utils/log"
	"agent/src/workflow"

	"github.com/gorilla/websocket"
	"golang.org/x/time/rate"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源（仅限开发环境）
	},
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}


const (
	// Time allowed to write a message to the peer.
	writeWait = 15 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 3 * 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 10*1024*1024
)

type client struct {
	ctx    context.Context
	conn   *websocket.Conn
	outMsg chan []byte
	authenticated bool
	apikey string
	agent  *Agent
	fnCancel context.CancelFunc
	rate   *rate.Limiter
	
	wantResponseId uint64
	wantResponse   chan []byte
	lock           sync.Mutex
	once           sync.Once
}

func newClient(conn *websocket.Conn, a *Agent) *client {
	ctx, fnCancel := context.WithCancel(a.ctx)
	return &client{
		ctx   : ctx,
		conn  : conn,
		outMsg: make(chan []byte, 32),
		agent : a,
		fnCancel: fnCancel,
		rate: rate.NewLimiter(5, 1), // 每秒5次限制
	}
}
func (c *client) setApiKey(apikey string) { 
	c.apikey = apikey
	c.authenticated = true
}

func (c *client) run() {
	go c.readPump()
	go c.writePump()
}

func (c *client) readPump() {
	defer c.Close()

	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error { c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	done := c.ctx.Done()
	for {
		select {			
		case <-done:
			return
		default:
		}

		_, msg, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Error(fmt.Sprintf("error: %v", err))
			}
			break
		}

		if c.rate.Allow() {
			workflow.Handle(c, msg)
		} else {
			c.WriteJson(&message.ReponseError{
				Error: "maximum 5 requests per second",
			}, true)
		}

	}
}


func (c *client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer ticker.Stop()
	defer c.Close()

	done := c.ctx.Done()
	for {
		select {			
		case message, ok := <-c.outMsg:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		case <-done:
			return
		}
	}
}

func (c *client) Close() {
	once.Do(func(){
		c.fnCancel()
		c.agent.removeClient(c)
		c.conn.Close()
		log.Info("client close", "remote", c.conn.RemoteAddr().String(), "connected", c.agent.clients.Size())
	})
}

func (c *client) WriteJson(data interface{}, block bool) error {
	b, err := json.MarshalIndent(data, "", "")
	if err != nil {
		return err
	}
	if block {
		select{
			case <-c.ctx.Done():
				return errors.New("connection closed")
			case c.outMsg <- b:
				return nil
		}
	} else {
		select {
			case <-c.ctx.Done():
				return errors.New("connection closed")
			case c.outMsg <- b:
				return nil
			default:
				go func(){
					select{
						case <-c.ctx.Done():
						case c.outMsg <- b:
					}
				}()
		}
	}
	return nil
}

func (c *client) WriteText(msg string, block bool) error {
	if block {
		select {
			case <-c.ctx.Done():
				return errors.New("connection closed")
			case c.outMsg <- []byte(msg):
				return nil
		}
	} else {
		select {
			case c.outMsg <- []byte(msg):
			default:
		}
	}
	return nil
}

// send string, block util flush
func (c *client) SendText(msg string) error {
	c.conn.SetWriteDeadline(time.Now().Add(writeWait))
	w, err := c.conn.NextWriter(websocket.TextMessage)
	if err != nil {
		return err
	}
	_, err = w.Write([]byte(msg))
	if err != nil {
		return err
	}
	err = w.Close()
	return err
}

func (c *client) GetCtx() context.Context {
	return c.ctx
}

func (c *client) GetChannel() chan<- []byte {
	return c.outMsg
}


func (c *client) WantResponse(requestId uint64, ch chan []byte) {
	c.lock.Lock()
	c.wantResponseId, c.wantResponse = requestId, ch
	c.lock.Unlock()
	close(c.wantResponse)
}

func (c *client) NewWantResponse(requestId uint64, msg []byte) error {
	c.lock.Lock()
	defer c.lock.Unlock()
	if requestId != c.wantResponseId || c.wantResponse == nil {
		return errors.New("unexpected msg")
	}
	c.wantResponse <- msg
	return nil
}