package main

import (
	"bufio"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// ---------------------------------------------------------------------------
// Signal metadata (populated by DISCOVER)
// ---------------------------------------------------------------------------

type SignalMeta struct {
	Name       string   `json:"name"`
	ID         uint32   `json:"id"`
	Type       string   `json:"type"`
	Dimensions uint8    `json:"dimensions"`
	Elements   uint32   `json:"elements"`
	Names      []string // canonical + alias names mapping to this ID
}

// ---------------------------------------------------------------------------
// Outbound WS message helpers
// ---------------------------------------------------------------------------

type TelemetrySample struct {
	ID     uint32    `json:"id"`
	Names  []string  `json:"names"`
	Ts     float64   `json:"ts"`
	Values []float64 `json:"values"`
}

func broadcast(hub *Hub, v any) {
	b, err := json.Marshal(v)
	if err != nil {
		return
	}
	hub.broadcast(b)
}

// ---------------------------------------------------------------------------
// MarteClient
// ---------------------------------------------------------------------------

type MarteClient struct {
	hub *Hub

	mu      sync.Mutex
	tcpConn net.Conn
	writer  *bufio.Writer

	cmdMu sync.Mutex // serialise TCP writes

	sigMu   sync.RWMutex
	signals map[uint32]*SignalMeta // id -> meta

	baseTs    uint64
	baseTsSet bool
	basesMu   sync.Mutex

	connected int32 // atomic bool

	lastWriteMs int64 // atomic; updated on every TCP write, used by keepalive

	stopCh chan struct{}

	// accumulates signals across DISCOVER_PART chunks; merged on final DISCOVER
	discoverAcc []discoverSignalJSON

	// cached last-known ports (updated by SERVICE_INFO)
	host    string
	cmdPort int
	udpPort int
	logPort int
}

func newMarteClient(hub *Hub) *MarteClient {
	return &MarteClient{
		hub:     hub,
		signals: make(map[uint32]*SignalMeta),
		stopCh:  make(chan struct{}),
	}
}

func (m *MarteClient) IsConnected() bool {
	return atomic.LoadInt32(&m.connected) == 1
}

// SendCachedState sends the current signal list to a freshly connected browser.
func (m *MarteClient) SendCachedState(c *WSClient) {
	m.sigMu.RLock()
	defer m.sigMu.RUnlock()
	if len(m.signals) == 0 {
		return
	}
	sigs := make([]*SignalMeta, 0, len(m.signals))
	for _, s := range m.signals {
		sigs = append(sigs, s)
	}
	b, _ := json.Marshal(map[string]any{"type": "signal_cache", "signals": sigs})
	select {
	case c.send <- b:
	default:
	}
}

// ---------------------------------------------------------------------------
// Connect / Disconnect
// ---------------------------------------------------------------------------

func (m *MarteClient) Connect(host string, cmdPort, udpPort, logPort int) {
	m.Disconnect()

	m.mu.Lock()
	m.host = host
	m.cmdPort = cmdPort
	m.udpPort = udpPort
	m.logPort = logPort
	m.stopCh = make(chan struct{})
	m.mu.Unlock()

	go m.runTCP(host, cmdPort)
	go m.runUDP(host, udpPort)
	go m.runLog(host, logPort)
}

func (m *MarteClient) Disconnect() {
	m.mu.Lock()
	select {
	case <-m.stopCh:
		// already closed
	default:
		close(m.stopCh)
	}
	if m.tcpConn != nil {
		m.tcpConn.Close()
		m.tcpConn = nil
	}
	m.mu.Unlock()

	atomic.StoreInt32(&m.connected, 0)
	m.basesMu.Lock()
	m.baseTsSet = false
	m.basesMu.Unlock()
	m.discoverAcc = nil
}

func (m *MarteClient) stopped() bool {
	select {
	case <-m.stopCh:
		return true
	default:
		return false
	}
}

// ---------------------------------------------------------------------------
// TCP command channel
// ---------------------------------------------------------------------------

func (m *MarteClient) runTCP(host string, port int) {
	addr := fmt.Sprintf("%s:%d", host, port)
	for !m.stopped() {
		conn, err := net.DialTimeout("tcp", addr, 5*time.Second)
		if err != nil {
			time.Sleep(2 * time.Second)
			continue
		}
		conn.(*net.TCPConn).SetNoDelay(true)

		m.mu.Lock()
		m.tcpConn = conn
		m.writer = bufio.NewWriter(conn)
		m.mu.Unlock()

		atomic.StoreInt32(&m.connected, 1)
		broadcast(m.hub, map[string]any{"type": "connected"})

		// Send SERVICE_INFO to auto-discover ports
		m.writeCmd("SERVICE_INFO")

		go m.runKeepalive()

		m.readLoop(conn)

		atomic.StoreInt32(&m.connected, 0)
		broadcast(m.hub, map[string]any{"type": "disconnected"})

		m.mu.Lock()
		m.tcpConn = nil
		m.writer = nil
		m.mu.Unlock()

		if !m.stopped() {
			time.Sleep(2 * time.Second)
		}
	}
}

func (m *MarteClient) writeCmd(cmd string) {
	m.cmdMu.Lock()
	defer m.cmdMu.Unlock()
	m.mu.Lock()
	w := m.writer
	m.mu.Unlock()
	if w == nil {
		return
	}
	log.Printf("[→MARTe] %s", cmd)
	broadcast(m.hub, map[string]any{
		"type": "log", "time": time.Now().Format("15:04:05.000"),
		"level": "CMD", "message": fmt.Sprintf("→ %s", cmd),
	})
	w.WriteString(cmd + "\n")
	w.Flush()
	atomic.StoreInt64(&m.lastWriteMs, time.Now().UnixMilli())
}

// runKeepalive sends INFO every 20 s when no other command has been written
// in the last 20 s, keeping the DebugService 30 s idle timer from firing.
func (m *MarteClient) runKeepalive() {
	ticker := time.NewTicker(20 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-m.stopCh:
			return
		case <-ticker.C:
			if !m.IsConnected() {
				continue
			}
			idleMs := time.Now().UnixMilli() - atomic.LoadInt64(&m.lastWriteMs)
			if idleMs >= 20_000 {
				m.writeCmd("INFO")
			}
		}
	}
}

func (m *MarteClient) SendCommand(cmd string) {
	m.writeCmd(cmd)
}

func (m *MarteClient) readLoop(conn net.Conn) {
	scanner := bufio.NewScanner(conn)
	scanner.Buffer(make([]byte, 8*1024*1024), 8*1024*1024)

	var jsonAcc strings.Builder
	inJSON := false

	for scanner.Scan() {
		if m.stopped() {
			return
		}
		line := scanner.Text()
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}

		// Detect start of JSON block
		if !inJSON && strings.HasPrefix(trimmed, "{") {
			inJSON = true
			jsonAcc.Reset()
		}

		if inJSON {
			jsonAcc.WriteString(trimmed)
			tag, done := detectJSONDone(trimmed)
			if done {
				inJSON = false
				raw := jsonAcc.String()
				// Strip trailing sentinel
				idx := strings.Index(raw, "OK "+tag)
				if idx >= 0 {
					raw = strings.TrimSpace(raw[:idx])
				}
				m.handleJSONResponse(tag, raw)
				jsonAcc.Reset()
			}
		} else {
			m.handleTextLine(trimmed)
		}
	}
}

// detectJSONDone returns (tag, true) when the line is exactly "OK <TAG>".
// DISCOVER_PART must appear before DISCOVER so partial chunks are not mistaken
// for the final chunk (though with exact matching this is not strictly
// necessary; kept for clarity).
func detectJSONDone(line string) (string, bool) {
	tags := []string{
		"DISCOVER_PART", "DISCOVER", "TREE", "INFO", "CONFIG", "STEP_STATUS",
		"VALUE", "MSG", "TRACE", "FORCE", "UNFORCE", "BREAK",
		"PAUSE", "RESUME", "STEP", "MONITOR", "UNMONITOR", "LS",
	}
	for _, t := range tags {
		if line == "OK "+t {
			return t, true
		}
	}
	return "", false
}

func (m *MarteClient) handleJSONResponse(tag, data string) {
	log.Printf("[←MARTe] %s  %d bytes", tag, len(data))
	broadcast(m.hub, map[string]any{
		"type": "log", "time": time.Now().Format("15:04:05.000"),
		"level": "RESP", "message": fmt.Sprintf("← %s (%d B)", tag, len(data)),
	})
	switch tag {
	case "DISCOVER_PART":
		// Accumulate this chunk; do NOT broadcast — we wait for the final chunk.
		var resp discoverResp
		if err := json.Unmarshal([]byte(data), &resp); err != nil {
			log.Printf("[DISCOVER_PART] parse error: %v", err)
			return
		}
		m.discoverAcc = append(m.discoverAcc, resp.Signals...)
		log.Printf("[DISCOVER_PART] accumulated %d signals (total so far: %d)",
			len(resp.Signals), len(m.discoverAcc))
		return

	case "DISCOVER":
		// Merge any previously accumulated part-chunks with this final chunk.
		var resp discoverResp
		if err := json.Unmarshal([]byte(data), &resp); err != nil {
			log.Printf("[DISCOVER] parse error: %v", err)
			return
		}
		all := append(m.discoverAcc, resp.Signals...)
		m.discoverAcc = nil
		m.parseDiscoverSignals(all)
		// Re-marshal the merged list so the browser gets a single consistent blob.
		merged, _ := json.Marshal(discoverResp{Signals: all})
		broadcast(m.hub, map[string]any{
			"type": "response", "tag": "DISCOVER", "data": string(merged),
		})
		return

	case "TREE":
		// Parse server-side and stream pre-digested flat-node batches to the
		// browser so it never has to JSON.parse a multi-MB string.
		go m.sendTreeBatches(data)
		return
	}
	broadcast(m.hub, map[string]any{
		"type": "response",
		"tag":  tag,
		"data": data,
	})
}

// ---------------------------------------------------------------------------
// TREE streaming
// ---------------------------------------------------------------------------

// TreeNode mirrors the DebugService TREE JSON structure.
type TreeNode struct {
	Name        string      `json:"Name"`
	Class       string      `json:"Class"`
	IsTraceable bool        `json:"IsTraceable"`
	IsForcable  bool        `json:"IsForcable"`
	Elements    uint32      `json:"Elements"`
	Children    []*TreeNode `json:"Children"`
}

// FlatNode is a pre-digested, compact representation sent to the browser.
type FlatNode struct {
	Path   string `json:"p"`
	Name   string `json:"n"`
	Class  string `json:"c"`
	Parent string `json:"par"`
	Tr     bool   `json:"tr,omitempty"`
	Fo     bool   `json:"fo,omitempty"`
	El     uint32 `json:"el,omitempty"`
	HasCh  bool   `json:"ch,omitempty"`
}

func flattenTree(node *TreeNode, parentPath string, out *[]FlatNode) {
	path := node.Name
	if parentPath != "" {
		path = parentPath + "." + node.Name
	}
	hasCh := len(node.Children) > 0
	*out = append(*out, FlatNode{
		Path:  path,
		Name:  node.Name,
		Class: node.Class,
		Parent: parentPath,
		Tr:    node.IsTraceable,
		Fo:    node.IsForcable,
		El:    node.Elements,
		HasCh: hasCh,
	})
	for _, c := range node.Children {
		flattenTree(c, path, out)
	}
}

// sendTreeBatches parses the raw TREE JSON and streams compact FlatNode
// batches to all browsers. Each browser processes one batch per animation
// frame, keeping the UI responsive during large tree loads.
func (m *MarteClient) sendTreeBatches(raw string) {
	var root TreeNode
	if err := json.Unmarshal([]byte(raw), &root); err != nil {
		log.Printf("[TREE] parse error: %v — falling back to raw JSON", err)
		broadcast(m.hub, map[string]any{"type": "response", "tag": "TREE", "data": raw})
		return
	}

	var nodes []FlatNode
	if root.Name == "Root" && len(root.Children) > 0 {
		for _, c := range root.Children {
			flattenTree(c, "", &nodes)
		}
	} else {
		flattenTree(&root, "", &nodes)
	}

	total := len(nodes)
	log.Printf("[TREE] streaming %d flat nodes to browser", total)
	broadcast(m.hub, map[string]any{"type": "tree_clear", "total": total})

	const batchSize = 50
	for i := 0; i < total; i += batchSize {
		end := i + batchSize
		if end > total {
			end = total
		}
		broadcast(m.hub, map[string]any{
			"type":  "tree_batch",
			"nodes": nodes[i:end],
			"done":  end == total,
		})
	}
}

func (m *MarteClient) handleTextLine(line string) {
	if strings.HasPrefix(line, "OK SERVICE_INFO") {
		// OK SERVICE_INFO TCP_CTRL:8110 UDP_STREAM:8111 TCP_LOG:8082 STATE:RUNNING
		parts := strings.Fields(line)
		newUDP, newLog := 0, 0
		for _, p := range parts {
			if strings.HasPrefix(p, "UDP_STREAM:") {
				fmt.Sscanf(p[11:], "%d", &newUDP)
			}
			if strings.HasPrefix(p, "TCP_LOG:") {
				fmt.Sscanf(p[8:], "%d", &newLog)
			}
		}
		// Always forward as a response so the browser can display it.
		broadcast(m.hub, map[string]any{
			"type": "response",
			"tag":  "SERVICE_INFO",
			"data": line[len("OK SERVICE_INFO "):],
		})
		if newUDP > 0 || newLog > 0 {
			broadcast(m.hub, map[string]any{
				"type":     "service_config",
				"udp_port": newUDP,
				"log_port": newLog,
			})
			// Restart UDP/log workers with updated ports if they differ
			m.mu.Lock()
			host := m.host
			oldUDP := m.udpPort
			oldLog := m.logPort
			if newUDP > 0 {
				m.udpPort = newUDP
			}
			if newLog > 0 {
				m.logPort = newLog
			}
			m.mu.Unlock()
			if newUDP > 0 && newUDP != oldUDP {
				go m.runUDP(host, newUDP)
			}
			if newLog > 0 && newLog != oldLog {
				go m.runLog(host, newLog)
			}
		}
	}
	// Forward all text lines to browsers.
	broadcast(m.hub, map[string]any{
		"type": "text_line",
		"data": line,
	})
}

// ---------------------------------------------------------------------------
// DISCOVER parsing
// ---------------------------------------------------------------------------

type discoverSignalJSON struct {
	Name       string `json:"name"`
	ID         uint32 `json:"id"`
	Type       string `json:"type"`
	Dimensions uint8  `json:"dimensions"`
	Elements   uint32 `json:"elements"`
}

type discoverResp struct {
	Signals []discoverSignalJSON `json:"Signals"`
}

func (m *MarteClient) parseDiscoverSignals(sigs []discoverSignalJSON) {
	m.sigMu.Lock()
	defer m.sigMu.Unlock()
	m.signals = make(map[uint32]*SignalMeta, len(sigs))
	for _, s := range sigs {
		el := s.Elements
		if el == 0 {
			el = 1
		}
		// Multiple DISCOVER entries can share the same ID (canonical DataSource
		// path + one GAM alias per broker).  Merge them into one SignalMeta so
		// that Names carries every alias.  This is required for the telemetry
		// receiver to match whichever name the user traced.
		if existing, ok := m.signals[s.ID]; ok {
			existing.Names = append(existing.Names, s.Name)
			continue
		}
		meta := &SignalMeta{
			Name:       s.Name,
			ID:         s.ID,
			Type:       s.Type,
			Dimensions: s.Dimensions,
			Elements:   el,
			Names:      []string{s.Name},
		}
		m.signals[s.ID] = meta
	}
	log.Printf("[DISCOVER] registered %d unique signals from %d entries", len(m.signals), len(sigs))
}

// ---------------------------------------------------------------------------
// UDP telemetry
// ---------------------------------------------------------------------------

const udpMagic uint32 = 0xDA7A57AD

func (m *MarteClient) runUDP(host string, port int) {
	addr := fmt.Sprintf("0.0.0.0:%d", port)
	udpAddr, err := net.ResolveUDPAddr("udp4", addr)
	if err != nil {
		return
	}
	conn, err := net.ListenUDP("udp4", udpAddr)
	if err != nil {
		return
	}
	defer conn.Close()
	conn.SetReadBuffer(10 * 1024 * 1024)
	conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))

	buf := make([]byte, 65535)
	var lastSeq uint32
	var hasLastSeq bool
	var totalPkts uint64
	var dropped uint64

	for !m.stopped() {
		conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
		n, _, err := conn.ReadFromUDP(buf)
		if err != nil {
			continue
		}
		if n < 20 {
			continue
		}
		if binary.LittleEndian.Uint32(buf[0:4]) != udpMagic {
			continue
		}

		seq := binary.LittleEndian.Uint32(buf[4:8])
		if hasLastSeq && seq > lastSeq+1 {
			dropped += uint64(seq - lastSeq - 1)
		}
		lastSeq = seq
		hasLastSeq = true

		totalPkts++
		if totalPkts%500 == 0 {
			broadcast(m.hub, map[string]any{
				"type":    "udp_stats",
				"packets": totalPkts,
				"dropped": dropped,
			})
		}

		count := binary.LittleEndian.Uint32(buf[16:20])
		offset := 20

		// Resolve base timestamp once per packet.
		m.basesMu.Lock()
		if !m.baseTsSet && n >= 20+12 {
			m.baseTs = binary.LittleEndian.Uint64(buf[20+4 : 20+12])
			m.baseTsSet = true
		}
		baseTs := m.baseTs
		baseTsSet := m.baseTsSet
		m.basesMu.Unlock()

		m.sigMu.RLock()
		samples := make([]TelemetrySample, 0, count)
		for i := uint32(0); i < count; i++ {
			if offset+16 > n {
				break
			}
			id := binary.LittleEndian.Uint32(buf[offset : offset+4])
			tsRaw := binary.LittleEndian.Uint64(buf[offset+4 : offset+12])
			size := binary.LittleEndian.Uint32(buf[offset+12 : offset+16])
			offset += 16

			if offset+int(size) > n {
				break
			}
			dataSlice := buf[offset : offset+int(size)]
			offset += int(size)

			tsS := 0.0
			if baseTsSet && tsRaw >= baseTs {
				tsS = float64(tsRaw-baseTs) / 1_000_000.0
			}

			meta, ok := m.signals[id]
			if !ok {
				continue
			}

			elements := meta.Elements
			if elements == 0 {
				elements = 1
			}
			typeSize := uint32(size)
			if elements > 0 {
				typeSize = uint32(size) / elements
			}
			vals := make([]float64, 0, elements)
			t := meta.Type

			for e := uint32(0); e < elements; e++ {
				off := e * typeSize
				if off+typeSize > uint32(len(dataSlice)) {
					break
				}
				elem := dataSlice[off : off+typeSize]
				vals = append(vals, parseTypedValue(elem, typeSize, t))
			}

			samples = append(samples, TelemetrySample{
				ID:     id,
				Names:  meta.Names,
				Ts:     tsS,
				Values: vals,
			})
		}
		m.sigMu.RUnlock()

		if len(samples) > 0 {
			broadcast(m.hub, map[string]any{
				"type":    "telemetry",
				"seq":     seq,
				"signals": samples,
			})
		}
	}
}

func parseTypedValue(data []byte, size uint32, t string) float64 {
	switch size {
	case 1:
		if strings.Contains(t, "u") {
			return float64(data[0])
		}
		return float64(int8(data[0]))
	case 2:
		if len(data) < 2 {
			return 0
		}
		v := binary.LittleEndian.Uint16(data[:2])
		if strings.Contains(t, "u") {
			return float64(v)
		}
		return float64(int16(v))
	case 4:
		if len(data) < 4 {
			return 0
		}
		v := binary.LittleEndian.Uint32(data[:4])
		if strings.Contains(t, "float") || strings.Contains(t, "float32") {
			return float64(math.Float32frombits(v))
		}
		if strings.Contains(t, "u") {
			return float64(v)
		}
		return float64(int32(v))
	case 8:
		if len(data) < 8 {
			return 0
		}
		v := binary.LittleEndian.Uint64(data[:8])
		if strings.Contains(t, "float") || strings.Contains(t, "float64") || strings.Contains(t, "double") {
			return math.Float64frombits(v)
		}
		if strings.Contains(t, "u") {
			return float64(v)
		}
		return float64(int64(v))
	}
	return 0
}

// ---------------------------------------------------------------------------
// TCP log channel
// ---------------------------------------------------------------------------

func (m *MarteClient) runLog(host string, port int) {
	addr := fmt.Sprintf("%s:%d", host, port)
	for !m.stopped() {
		conn, err := net.DialTimeout("tcp", addr, 5*time.Second)
		if err != nil {
			time.Sleep(2 * time.Second)
			continue
		}
		scanner := bufio.NewScanner(conn)
		for scanner.Scan() {
			if m.stopped() {
				conn.Close()
				return
			}
			line := strings.TrimSpace(scanner.Text())
			if strings.HasPrefix(line, "LOG ") {
				rest := line[4:]
				idx := strings.Index(rest, " ")
				if idx < 0 {
					continue
				}
				level := rest[:idx]
				msg := rest[idx+1:]
				broadcast(m.hub, map[string]any{
					"type":    "log",
					"time":    time.Now().Format("15:04:05.000"),
					"level":   level,
					"message": msg,
				})
			}
		}
		conn.Close()
		if !m.stopped() {
			time.Sleep(2 * time.Second)
		}
	}
}