// main.go
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"

	"github.com/mbndr/figlet4go"
	"nhooyr.io/websocket"

	// "nhooyr.io/websocket/wsjson"
	"os"
	"strconv"
	"sync"
)

type InitPayload struct {
	RoomID    string `json:"roomId"`
	IsBackend bool   `json:"isBackend"`
}

type Config struct {
	Port int `json:"port"`
}

func LoadConfig(path string) Config {
	file, err := os.Open(path)
	if err != nil {
		log.Fatalf("Failed to open config file: %v", err)
	}
	defer file.Close()

	decoder := json.NewDecoder(file)
	var cfg Config
	if err := decoder.Decode(&cfg); err != nil {
		log.Fatalf("Failed to parse config file: %v", err)
	}
	return cfg
}

type Client struct {
	conn      *websocket.Conn
	roomID    string
	isBackend bool
	send      chan []byte
	ip        string
}

type Hub struct {
	rooms      map[string]map[*Client]bool
	backend    map[string]*Client
	register   chan *Client
	unregister chan *Client
	broadcast  chan Message
	mu         sync.Mutex
}

type Message struct {
	RoomID string
	Data   []byte
}

func NewHub() *Hub {
	return &Hub{
		rooms:      make(map[string]map[*Client]bool),
		backend:    make(map[string]*Client),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		broadcast:  make(chan Message),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			if client.isBackend {
				if _, exists := h.backend[client.roomID]; exists {
					log.Printf("[%s] RoomId Exists", client.ip)
					client.roomID = ""
					err := client.conn.Close(websocket.StatusPolicyViolation, "Backend already exists")
					if err != nil {
						log.Printf("[%s] Error closing connection", client.ip)
					}
					h.mu.Unlock()
					continue
				}
				log.Printf("[%s] Registering backend for room %s", client.ip, client.roomID)
				h.backend[client.roomID] = client
			} else {
				if _, ok := h.backend[client.roomID]; !ok {
					log.Printf("[%s] RoomId Not Exists: %s", client.ip, client.roomID)
					err := client.conn.Close(websocket.StatusPolicyViolation, "No backend")
					if err != nil {
						log.Printf("[%s] Error closing connection", client.ip)
					}
					h.mu.Unlock()
					continue
				}
				log.Printf("[%s] Registering client for room %s", client.ip, client.roomID)
			}
			if _, ok := h.rooms[client.roomID]; !ok {
				h.rooms[client.roomID] = make(map[*Client]bool)
			}
			h.rooms[client.roomID][client] = true
			h.mu.Unlock()
		case client := <-h.unregister:
			h.mu.Lock()
			if room, ok := h.rooms[client.roomID]; ok {
				// Only close/delete if the client actually exists in the Hub's map.
				// It might have been removed by broadcast already.
				if _, exists := room[client]; exists {
					delete(room, client)
					close(client.send)
				}

				// Cleanup Backends as before
				if client.isBackend {
					delete(h.backend, client.roomID)
					// Note: If we iterate 'room' here, we are closing connections for Viewer clients.
					// This is 'destroy room' logic.
					// Need to be careful not to double-close their channels too if they are unregistered concurrently.
					// The safe way is to send them a shutdown signal or close their conn, letting readPump trigger unregister.
					// But here we are forcing closure.
					for c := range room {
						log.Printf("[%s] Closing connection for client in room %s", c.ip, client.roomID)
						// We don't close c.send here, we close the socket.
						// c.readPump will exit and likely trigger unregister for 'c'.
						err := c.conn.Close(websocket.StatusNormalClosure, "Backend disconnected")
						if err != nil {
							log.Printf("[%s] Error closing connection", client.ip)
						}
					}
				}
				if len(room) == 0 {
					log.Printf("[%s] Unregistering backend for room %s", client.ip, client.roomID)
					delete(h.rooms, client.roomID)
				}
			}
			h.mu.Unlock()
		case msg := <-h.broadcast:
			h.mu.Lock()
			if clients, ok := h.rooms[msg.RoomID]; ok {
				for c := range clients {
					// Don't send back to self (optional, but good practice)
					// We don't have 'sender' info here easily unless we add it to Message struct.
					// For now, checking `isBackend` is impossible since we removed it.
					// Let's blindly broadcast to all for simplicity, or we update Message to include Sender.

					// Simple broadcast to all in room
					select {
					case c.send <- msg.Data:
					default:
						close(c.send)
						delete(clients, c)
					}
				}
			}
			h.mu.Unlock()
		}
	}
}

func serveWebSocket(hub *Hub, w http.ResponseWriter, r *http.Request) {
	// conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
	// 	CompressionMode: websocket.CompressionContextTakeover,
	// 	OriginPatterns:  []string{"*"},
	// })

	// 【修改点】使用混合配置
	cOps := &websocket.AcceptOptions{
		// 1. 必须禁用压缩 (为了解决 iOS 连接断开/闪断问题)
		// iOS WebKit 对压缩极其敏感，必须关掉
		CompressionMode: websocket.CompressionDisabled,

		// 2. 恢复使用 OriginPatterns (为了解决 "无法连接" 问题)
		// 既然你之前的代码用这个能连上 C++ 和安卓，我们就改回这个。
		// * 号表示允许任何域名连接
		OriginPatterns: []string{"*"},
	}

	// 发起握手
	conn, err := websocket.Accept(w, r, cOps)

	if err != nil {
		log.Println("Upgrade error:", err)
		return
	}
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()

	client := &Client{
		conn: conn,
		ip:   r.RemoteAddr,
		send: make(chan []byte, 256),
	}
	go client.writePump(ctx)
	client.readPump(ctx, hub)
}

func (c *Client) readPump(ctx context.Context, hub *Hub) {
	defer func() {
		hub.unregister <- c
		err := c.conn.Close(websocket.StatusNormalClosure, "Disconnected")
		if err != nil {
			log.Printf("[%s] Error closing connection", c.ip)
		}
	}()

	// Read raw message to debug content
	_, msg, err := c.conn.Read(ctx)
	if err != nil {
		log.Printf("[%s] Read handshake error: %v", c.ip, err)
		return
	}

	log.Printf("[%s] Received Handshake Raw: %s", c.ip, string(msg))

	var payload InitPayload
	if err := json.Unmarshal(msg, &payload); err != nil {
		log.Printf("[%s] Handshake JSON parse error: %v", c.ip, err)
		return
	}

	c.roomID = payload.RoomID
	c.isBackend = payload.IsBackend
	hub.register <- c

	// Send response manually since we used raw Read
	response, _ := json.Marshal(map[string]string{"status": "ok"})
	err = c.conn.Write(ctx, websocket.MessageText, response)
	if err != nil {
		log.Printf("[%s] Error writing response: %v", c.ip, err)
	}

	for {
		_, msg, err := c.conn.Read(ctx)
		if err != nil {
			log.Printf("[%s] Read message error: %v", c.ip, err)
			return
		}

		// Debug Log: Print EVERYTHING received
		// log.Printf("[%s] Received Frame: Type=%v, Size=%d, IsBackend=%v", c.ip, typ, len(msg), c.isBackend)

		if c.isBackend {
			// Forward both Binary (2) and Text (1)
			// Note: If C++ sends invalid UTF-8 as Text, Read() above might have already errored.
			hub.broadcast <- Message{RoomID: c.roomID, Data: msg}
		}
	}
}

func (c *Client) writePump(ctx context.Context) {
	for {
		select {
		case msg, ok := <-c.send:
			if !ok {
				return
			}
			err := c.conn.Write(ctx, websocket.MessageBinary, msg)
			if err != nil {
				log.Printf("[%s] Error writing message: %v", c.ip, err)
			}
		case <-ctx.Done():
			return
		}
	}
}

var config Config

func ShowFiglet(str string) {
	ascii := figlet4go.NewAsciiRender()
	options := figlet4go.NewRenderOptions()
	options.FontName = "larry3d"
	renderStr, _ := ascii.RenderOpts(str, options)
	fmt.Print(renderStr)
}

func main() {

	ShowFiglet("Transfer")

	config = LoadConfig("config.json")

	hub := NewHub()
	go hub.Run()

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		serveWebSocket(hub, w, r)
	})

	addr := ":" + strconv.Itoa(config.Port)
	log.Println("WebSocket relay server started on", addr)
	log.Fatal(http.ListenAndServe(addr, nil))
}
