package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gdamore/tcell/v2"
	"github.com/rivo/tview"
)

// --- Models ---

type Signal struct {
	Name string `json:"name"`
	ID   uint32 `json:"id"`
	Type string `json:"type"`
}

type DiscoverResponse struct {
	Signals []Signal `json:"Signals"`
}

type TraceValue struct {
	Value      string
	LastUpdate time.Time
}

type AppState struct {
	Signals       map[string]Signal
	IDToSignal    map[uint32]string
	ForcedSignals map[string]string
	TracedSignals map[string]TraceValue
	mu            sync.RWMutex
}

var state = AppState{
	Signals:       make(map[string]Signal),
	IDToSignal:    make(map[uint32]string),
	ForcedSignals: make(map[string]string),
	TracedSignals: make(map[string]TraceValue),
}

// --- UI Components ---

var (
	app              *tview.Application
	forcedPane       *tview.Table
	tracedPane       *tview.Table
	statusPane       *tview.TextView
	frameworkLogPane *tview.TextView
	cmdInput         *tview.InputField
	mainFlex         *tview.Flex
)

// logToStatus prints command output WITHOUT timestamp
func logToStatus(msg string) {
	fmt.Fprintf(statusPane, "%s\n", msg)
}

// logToFramework prints framework logs WITH timestamp
func logToFramework(level, msg string) {
	color := "white"
	switch level {
	case "FatalError", "OSError", "ParametersError":
		color = "red"
	case "Warning":
		color = "yellow"
	case "Information":
		color = "green"
	case "Debug":
		color = "blue"
	}
	fmt.Fprintf(frameworkLogPane, "[%s] [[%s]%s[-]] %s\n", time.Now().Format("15:04:05.000"), color, level, msg)
}

// --- Network Logic ---

var (
	tcpConn      net.Conn
	tcpMu        sync.Mutex
	responseChan = make(chan string, 2000)
	connected    = false
)

func connectAndMonitor(addr string) {
	for {
		conn, err := net.Dial("tcp", addr)
		if err != nil {
			app.QueueUpdateDraw(func() { statusPane.SetText("[red]Connection failed, retrying...\n") })
			connected = false
			time.Sleep(2 * time.Second)
			continue
		}

		if tc, ok := conn.(*net.TCPConn); ok {
			tc.SetNoDelay(true)
		}

		tcpConn = conn
		connected = true
		app.QueueUpdateDraw(func() { logToStatus("[green]Connected to DebugService (TCP 8080)") })

		go discover()

		buf := make([]byte, 4096)
		var line strings.Builder
		for {
			n, err := conn.Read(buf)
			if err != nil {
				app.QueueUpdateDraw(func() { logToStatus("[red]TCP Disconnected: " + err.Error()) })
				connected = false
				break
			}

			for i := 0; i < n; i++ {
				if buf[i] == '\n' {
					fullLine := strings.TrimSpace(line.String())
					line.Reset()

					if strings.HasPrefix(fullLine, "LOG ") {
						parts := strings.SplitN(fullLine[4:], " ", 2)
						if len(parts) == 2 {
							app.QueueUpdateDraw(func() { logToFramework(parts[0], parts[1]) })
						}
					} else if fullLine != "" {
						select {
						case responseChan <- fullLine:
						default:
						}
					}
				} else {
					line.WriteByte(buf[i])
				}
			}
		}
		tcpConn.Close()
		time.Sleep(2 * time.Second)
	}
}

func startUDPListener(port int) {
	addr, _ := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", port))
	udpConn, err := net.ListenUDP("udp", addr)
	if err != nil {
		app.QueueUpdateDraw(func() { logToStatus(fmt.Sprintf("UDP Error: %v", err)) })
		return
	}
	defer udpConn.Close()

	buf := make([]byte, 2048)
	for {
		n, _, err := udpConn.ReadFromUDP(buf)
		if err != nil || n < 20 {
			continue
		}

		magic := binary.LittleEndian.Uint32(buf[0:4])
		if magic != 0xDA7A57AD {
			continue
		}

		count := binary.LittleEndian.Uint32(buf[16:20])
		offset := 20
		now := time.Now()
		for i := uint32(0); i < count; i++ {
			if offset+8 > n {
				break
			}
			id := binary.LittleEndian.Uint32(buf[offset : offset+4])
			size := binary.LittleEndian.Uint32(buf[offset+4 : offset+8])
			offset += 8
			if offset+int(size) > n {
				break
			}
			data := buf[offset : offset+int(size)]
			offset += int(size)

			state.mu.Lock()
			if name, ok := state.IDToSignal[id]; ok {
				valStr := ""
				if size == 4 {
					valStr = fmt.Sprintf("%d", binary.LittleEndian.Uint32(data))
				} else if size == 8 {
					valStr = fmt.Sprintf("%d", binary.LittleEndian.Uint64(data))
				} else {
					valStr = fmt.Sprintf("%X", data)
				}
				state.TracedSignals[name] = TraceValue{Value: valStr, LastUpdate: now}
			}
			state.mu.Unlock()
		}

		app.QueueUpdateDraw(func() { updateTracedPane() })
	}
}

func discover() {
	tcpMu.Lock()
	defer tcpMu.Unlock()
	if !connected {
		return
	}

	app.QueueUpdateDraw(func() { logToStatus("Discovering signals...") })

	// Drain old responses
	for len(responseChan) > 0 {
		<-responseChan
	}

	fmt.Fprintf(tcpConn, "DISCOVER\n")

	var jsonBlock strings.Builder
	timeout := time.After(5 * time.Second)
	for {
		select {
		case line := <-responseChan:
			if line == "OK DISCOVER" {
				goto parsed
			}
			jsonBlock.WriteString(line)
		case <-timeout:
			app.QueueUpdateDraw(func() { logToStatus("[red]Discovery Timeout") })
			return
		}
	}
parsed:
	var resp DiscoverResponse
	if err := json.Unmarshal([]byte(jsonBlock.String()), &resp); err == nil {
		state.mu.Lock()
		state.IDToSignal = make(map[uint32]string)
		for _, s := range resp.Signals {
			state.Signals[s.Name] = s
			state.IDToSignal[s.ID] = s.Name
		}
		state.mu.Unlock()
		app.QueueUpdateDraw(func() { logToStatus(fmt.Sprintf("Discovered %d signals", len(resp.Signals))) })
	}
}

// --- UI Updates ---

func updateForcedPane() {
	forcedPane.Clear()
	forcedPane.SetCell(0, 0, tview.NewTableCell("Signal").SetTextColor(tcell.ColorYellow).SetAttributes(tcell.AttrBold))
	forcedPane.SetCell(0, 1, tview.NewTableCell("Value").SetTextColor(tcell.ColorYellow).SetAttributes(tcell.AttrBold))
	state.mu.RLock()
	defer state.mu.RUnlock()
	keys := make([]string, 0, len(state.ForcedSignals))
	for k := range state.ForcedSignals {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for i, k := range keys {
		forcedPane.SetCell(i+1, 0, tview.NewTableCell(k))
		forcedPane.SetCell(i+1, 1, tview.NewTableCell(state.ForcedSignals[k]).SetTextColor(tcell.ColorGreen))
	}
}

func updateTracedPane() {
	tracedPane.Clear()
	tracedPane.SetCell(0, 0, tview.NewTableCell("Signal").SetTextColor(tcell.ColorBlue).SetAttributes(tcell.AttrBold))
	tracedPane.SetCell(0, 1, tview.NewTableCell("Value").SetTextColor(tcell.ColorBlue).SetAttributes(tcell.AttrBold))
	tracedPane.SetCell(0, 2, tview.NewTableCell("Last Update").SetTextColor(tcell.ColorBlue).SetAttributes(tcell.AttrBold))
	state.mu.RLock()
	defer state.mu.RUnlock()
	keys := make([]string, 0, len(state.TracedSignals))
	for k := range state.TracedSignals {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for i, k := range keys {
		tv := state.TracedSignals[k]
		tracedPane.SetCell(i+1, 0, tview.NewTableCell(k))
		tracedPane.SetCell(i+1, 1, tview.NewTableCell(tv.Value))
		tracedPane.SetCell(i+1, 2, tview.NewTableCell(tv.LastUpdate.Format("15:04:05.000")))
	}
}

// --- Command Handler ---

func handleInput(text string) {
	parts := strings.Fields(text)
	if len(parts) == 0 {
		return
	}
	cmd := strings.ToUpper(parts[0])
	args := parts[1:]
	if !connected && cmd != "EXIT" && cmd != "QUIT" && cmd != "CLEAR" {
		logToStatus("[red]Not connected to server")
		return
	}

	switch cmd {
	case "LS":
		path := ""
		if len(args) > 0 {
			path = args[0]
		}
		go func() {
			tcpMu.Lock()
			defer tcpMu.Unlock()
			for len(responseChan) > 0 {
				<-responseChan
			}
			fmt.Fprintf(tcpConn, "LS %s\n", path)
			app.QueueUpdateDraw(func() { logToStatus(fmt.Sprintf("Listing nodes under %s...", path)) })
			timeout := time.After(2 * time.Second)
			for {
				select {
				case line := <-responseChan:
					if line == "OK LS" {
						return
					}
					app.QueueUpdateDraw(func() { logToStatus(line) })
				case <-timeout:
					app.QueueUpdateDraw(func() { logToStatus("[red]LS Timeout") })
					return
				}
			}
		}()
	case "FORCE":
		if len(args) < 2 {
			return
		}
		path, val := args[0], args[1]
		go func() {
			tcpMu.Lock()
			defer tcpMu.Unlock()
			fmt.Fprintf(tcpConn, "FORCE %s %s\n", path, val)
			select {
			case resp := <-responseChan:
				app.QueueUpdateDraw(func() {
					logToStatus("Server: " + resp)
					state.mu.Lock()
					state.ForcedSignals[path] = val
					state.mu.Unlock()
					updateForcedPane()
				})
			case <-time.After(1 * time.Second):
				app.QueueUpdateDraw(func() { logToStatus("[red]Server Timeout") })
			}
		}()
	case "UNFORCE":
		if len(args) < 1 {
			return
		}
		path := args[0]
		go func() {
			tcpMu.Lock()
			defer tcpMu.Unlock()
			fmt.Fprintf(tcpConn, "UNFORCE %s\n", path)
			select {
			case resp := <-responseChan:
				app.QueueUpdateDraw(func() {
					logToStatus("Server: " + resp)
					state.mu.Lock()
					delete(state.ForcedSignals, path)
					state.mu.Unlock()
					updateForcedPane()
				})
			case <-time.After(1 * time.Second):
				app.QueueUpdateDraw(func() { logToStatus("[red]Server Timeout") })
			}
		}()
	case "TRACE":
		if len(args) < 1 {
			return
		}
		path := args[0]
		decim := "1"
		if len(args) > 1 {
			decim = args[1]
		}
		go func() {
			tcpMu.Lock()
			defer tcpMu.Unlock()
			fmt.Fprintf(tcpConn, "TRACE %s 1 %s\n", path, decim)
			select {
			case resp := <-responseChan:
				app.QueueUpdateDraw(func() {
					logToStatus("Server: " + resp)
					state.mu.Lock()
					state.TracedSignals[path] = TraceValue{Value: "...", LastUpdate: time.Now()}
					state.mu.Unlock()
					updateTracedPane()
				})
			case <-time.After(1 * time.Second):
				app.QueueUpdateDraw(func() { logToStatus("[red]Server Timeout") })
			}
		}()
	case "UNTRACE":
		if len(args) < 1 {
			return
		}
		path := args[0]
		go func() {
			tcpMu.Lock()
			defer tcpMu.Unlock()
			fmt.Fprintf(tcpConn, "UNTRACE %s\n", path)
			select {
			case resp := <-responseChan:
				app.QueueUpdateDraw(func() {
					logToStatus("Server: " + resp)
					state.mu.Lock()
					delete(state.TracedSignals, path)
					state.mu.Unlock()
					updateTracedPane()
				})
			case <-time.After(1 * time.Second):
				app.QueueUpdateDraw(func() { logToStatus("[red]Server Timeout") })
			}
		}()
	case "DISCOVER":
		go discover()
	case "CLEAR":
		statusPane.Clear()
		frameworkLogPane.Clear()
	case "EXIT", "QUIT":
		app.Stop()
	default:
		logToStatus("Unknown command: " + cmd)
	}
}

func main() {
	app = tview.NewApplication()

	forcedPane = tview.NewTable()
	forcedPane.SetBorders(true)
	forcedPane.SetTitle(" Forced Signals ")
	forcedPane.SetBorder(true)

	tracedPane = tview.NewTable()
	tracedPane.SetBorders(true)
	tracedPane.SetTitle(" Traced Signals (Live) ")
	tracedPane.SetBorder(true)

	statusPane = tview.NewTextView()
	statusPane.SetDynamicColors(true)
	statusPane.SetRegions(true)
	statusPane.SetWordWrap(true)
	statusPane.SetChangedFunc(func() {
		statusPane.ScrollToEnd()
		app.Draw()
	})
	statusPane.SetTitle(" CLI Status / Command Output ")
	statusPane.SetBorder(true)

	frameworkLogPane = tview.NewTextView()
	frameworkLogPane.SetDynamicColors(true)
	frameworkLogPane.SetRegions(true)
	frameworkLogPane.SetWordWrap(true)
	frameworkLogPane.SetChangedFunc(func() {
		frameworkLogPane.ScrollToEnd()
		app.Draw()
	})
	frameworkLogPane.SetTitle(" MARTe2 Framework Logs ")
	frameworkLogPane.SetBorder(true)

	cmdInput = tview.NewInputField()
	cmdInput.SetLabel("marte_debug> ")
	cmdInput.SetFieldWidth(0)
	cmdInput.SetDoneFunc(func(key tcell.Key) {
		if key == tcell.KeyEnter {
			text := cmdInput.GetText()
			if text != "" {
				handleInput(text)
				cmdInput.SetText("")
			}
		}
	})

	flexTop := tview.NewFlex()
	flexTop.SetDirection(tview.FlexColumn)
	flexTop.AddItem(forcedPane, 0, 1, false)
	flexTop.AddItem(tracedPane, 0, 1, false)
	flexTop.AddItem(statusPane, 0, 1, false)

	mainFlex = tview.NewFlex()
	mainFlex.SetDirection(tview.FlexRow)
	mainFlex.AddItem(flexTop, 0, 2, false)
	mainFlex.AddItem(frameworkLogPane, 0, 1, false)
	mainFlex.AddItem(cmdInput, 1, 0, true)

	go connectAndMonitor("127.0.0.1:8080")
	go startUDPListener(8081)

	if err := app.SetRoot(mainFlex, true).EnableMouse(true).Run(); err != nil {
		panic(err)
	}
}
