package client

import (
	"io"
	"log"
	"net"
	pb "net-tunnel/pkg/proto"
	"sync"
	"time"
)

type ConnectionState struct {
	conn        net.Conn
	connID      string
	config      *pb.ProxyConfig
	stream      pb.TunnelService_EstablishControlConnectionClient
	closed      bool
	closedMutex sync.Mutex
	closeOnce   sync.Once
}

func NewConnectionState(conn net.Conn, connID string, config *pb.ProxyConfig, stream pb.TunnelService_EstablishControlConnectionClient) *ConnectionState {
	return &ConnectionState{
		conn:        conn,
		connID:      connID,
		config:      config,
		stream:      stream,
		closed:      false,
		closedMutex: sync.Mutex{},
		closeOnce:   sync.Once{},
	}
}

func (c *ConnectionState) Close() {
	c.closeOnce.Do(func() {
		c.closedMutex.Lock()
		defer c.closedMutex.Unlock()

		if !c.closed {
			c.closed = true
			c.conn.Close()
			log.Printf("Connection %s closed", c.connID)
		}
	})
}

func (c *ConnectionState) IsClosed() bool {
	c.closedMutex.Lock()
	defer c.closedMutex.Unlock()
	return c.closed
}

func (c *ConnectionState) WriteData(data []byte) error {
	c.closedMutex.Lock()
	defer c.closedMutex.Unlock()

	if c.closed {
		return io.EOF
	}

	_, err := c.conn.Write(data)
	if err != nil {
		c.Close()
		return err
	}
	return nil
}

func (c *ConnectionState) StartReading() {
	go func() {
		buffer := make([]byte, 4096)
		defer c.Close()

		for {
			if c.IsClosed() {
				return
			}

			if c.config.Type == pb.ProxyType_TCP {
				if tcpConn, ok := c.conn.(*net.TCPConn); ok {
					_ = tcpConn.SetReadDeadline(time.Now().Add(time.Minute * 3))
				}
			}

			n, err := c.conn.Read(buffer)
			if err != nil {
				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
					continue
				}

				if err == io.EOF {
					log.Printf("Connection %s closed by remote", c.connID)
				}
				return
			}

			if n > 0 {
				if c.IsClosed() {
					return
				}
				if err := c.stream.Send(&pb.Message{
					Content: &pb.Message_ProxyData{
						ProxyData: &pb.ProxyData{
							ConnId:      c.connID,
							Data:        buffer[:n],
							ProxyConfig: c.config,
						},
					},
				}); err != nil {
					log.Printf("Failed to send proxy data: %v", err)
					c.Close()
					return
				}
			}
		}
	}()
}