package pva

import (
	"bufio"
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"log/slog"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/uopi/gopva/pvdata"
)

// ---- Channel/Request state machine -----------------------------------

type chanState int

const (
	chanPending chanState = iota // CREATE_CHANNEL sent
	chanCreated                  // server ack'd
	chanFailed                   // CREATE_CHANNEL failed
)

type serverChannel struct {
	pvName string
	cid    uint32 // client channel ID
	sid    uint32 // server channel ID (assigned on creation)
	state  chanState

	// pending callbacks waiting for channel creation
	onCreate []func(sid uint32, err error)
}

type monitorSub struct {
	ioid    uint32
	sid     uint32 // server channel ID — needed for pipeline ACKs and DESTROY
	cid     uint32
	desc    pvdata.FieldDesc // structure descriptor (received on INIT response)
	updates chan MonitorEvent
}

// MonitorEvent carries a decoded monitor update.
type MonitorEvent struct {
	// Changed is the set of top-level field indices that changed.
	Changed pvdata.BitSet
	// Value is the full decoded structure value.
	Value pvdata.StructValue
	// Err is non-nil if this is a terminal error event.
	Err error
}

// conn is a single PVA TCP connection to one server.
type conn struct {
	nc net.Conn
	br *bufio.Reader // shared reader — used by handshake then readLoop
	bw *bufio.Writer
	mu sync.Mutex // protects writes + state maps

	// auto-incrementing IDs
	nextCID  atomic.Uint32
	nextIOID atomic.Uint32

	// state
	chans    map[uint32]*serverChannel // keyed by client CID
	byPV     map[string]*serverChannel // keyed by PV name
	monitors map[uint32]*monitorSub    // keyed by IOID

	done chan struct{}
}

func dial(ctx context.Context, addr string) (*conn, error) {
	nc, err := (&net.Dialer{}).DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}
	c := &conn{
		nc:       nc,
		br:       bufio.NewReader(nc),
		bw:       bufio.NewWriter(nc),
		chans:    make(map[uint32]*serverChannel),
		byPV:     make(map[string]*serverChannel),
		monitors: make(map[uint32]*monitorSub),
		done:     make(chan struct{}),
	}
	if err := c.handshake(); err != nil {
		nc.Close()
		return nil, err
	}
	go c.readLoop()
	return c, nil
}

// ---- Handshake --------------------------------------------------------

// handshake completes the PVA connection setup.
//
// Protocol sequence (spec §4.2):
//  1. Client → server: SET_BYTE_ORDER control (announces little-endian)
//  2. Server → client: SET_BYTE_ORDER control (may be omitted by some servers)
//  3. Server → client: CONNECTION_VALIDATION (0x01) — server's auth options
//  4. Client → server: CONNECTION_VALIDATION (0x01) — client selects auth
//
// The client MUST NOT send step 4 before receiving step 3.
func (c *conn) handshake() error {
	// Step 1: announce our byte order.
	if _, err := c.nc.Write(BuildMessage(CtrlSetByteOrder, flagControl, nil)); err != nil {
		return fmt.Errorf("pva handshake: send byte-order: %w", err)
	}

	// Step 2+3: read server messages until CONNECTION_VALIDATION arrives.
	// Servers typically send SET_BYTE_ORDER first, then CONNECTION_VALIDATION.
	for {
		hdr, err := ReadHeader(c.br)
		if err != nil {
			return fmt.Errorf("pva handshake: read server message: %w", err)
		}
		payload := make([]byte, hdr.Size)
		if _, err := io.ReadFull(c.br, payload); err != nil {
			return fmt.Errorf("pva handshake: read server payload: %w", err)
		}
		if hdr.isControl() {
			// SET_BYTE_ORDER from server — we always use LE, ignore.
			continue
		}
		if hdr.Command == CmdConnectionValid {
			break // server is ready for our reply
		}
		// Ignore unexpected messages (e.g. BEACON on same port).
	}

	// Step 4: send CONNECTION_VALIDATION response.
	// payload: clientReceiveBufferSize(4) + clientIntrospectionRegistryMaxSize(4) + authNZ(str)
	var buf bytes.Buffer
	binary.Write(&buf, binary.LittleEndian, uint32(0x00800000)) // 8 MiB receive buffer
	binary.Write(&buf, binary.LittleEndian, uint32(0x00000000)) // introspection registry size
	buf.WriteByte(0)                                            // authNZ = "" (anonymous)
	if _, err := c.nc.Write(BuildMessage(CmdConnectionValid, flagApp, buf.Bytes())); err != nil {
		return fmt.Errorf("pva handshake: send connection_valid: %w", err)
	}
	return nil
}

// ---- Write helpers ----------------------------------------------------

func (c *conn) send(msg []byte) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	_, err := c.nc.Write(msg)
	return err
}

// ---- Channel creation -------------------------------------------------

// openChannel ensures a channel for pvName exists and calls cb when ready.
func (c *conn) openChannel(pvName string, cb func(sid uint32, err error)) {
	c.mu.Lock()
	if ch, ok := c.byPV[pvName]; ok {
		switch ch.state {
		case chanCreated:
			sid := ch.sid
			c.mu.Unlock()
			cb(sid, nil)
			return
		case chanFailed:
			c.mu.Unlock()
			cb(0, fmt.Errorf("pva: channel %q creation failed", pvName))
			return
		default:
			ch.onCreate = append(ch.onCreate, cb)
			c.mu.Unlock()
			return
		}
	}
	cid := c.nextCID.Add(1)
	ch := &serverChannel{pvName: pvName, cid: cid, state: chanPending, onCreate: []func(uint32, error){cb}}
	c.chans[cid] = ch
	c.byPV[pvName] = ch
	c.mu.Unlock()

	// send CREATE_CHANNEL
	// payload: count(2) + cid(4) + pvName(string)
	var payload bytes.Buffer
	binary.Write(&payload, binary.LittleEndian, uint16(1)) // channel count = 1
	binary.Write(&payload, binary.LittleEndian, cid)
	pvdata.WriteString(&payload, pvName)
	if err := c.send(BuildMessage(CmdCreateChannel, flagApp, payload.Bytes())); err != nil {
		c.mu.Lock()
		delete(c.chans, cid)
		delete(c.byPV, pvName)
		c.mu.Unlock()
		cb(0, err)
	}
}

// ---- Monitor ----------------------------------------------------------

// subscribe sends an EVENT_ADD equivalent (MONITOR INIT) for pvName.
func (c *conn) subscribe(ctx context.Context, pvName string, ch chan MonitorEvent) {
	c.openChannel(pvName, func(sid uint32, err error) {
		if err != nil {
			ch <- MonitorEvent{Err: err}
			return
		}
		ioid := c.nextIOID.Add(1)

		// Build MONITOR INIT payload:
		// sid(4) + ioid(4) + subCmd(1=INIT) + pvRequest(FieldDesc+Value)
		var payload bytes.Buffer
		binary.Write(&payload, binary.LittleEndian, sid)
		binary.Write(&payload, binary.LittleEndian, ioid)
		payload.WriteByte(SubCmdInit) // subCmd = INIT

		// pvRequest: empty structure (field "field" selecting all)
		// Encoded as: TypeCodeStruct + typeID("") + nfields(0)
		// This selects the whole top-level structure.
		payload.WriteByte(pvdata.TypeCodeStruct) // FieldDesc type
		pvdata.WriteString(&payload, "")         // typeID
		pvdata.WriteSize(&payload, 0)            // no sub-fields = select all

		c.mu.Lock()
		var cid uint32
		for _, sch := range c.chans {
			if sch.sid == sid {
				cid = sch.cid
				break
			}
		}
		ms := &monitorSub{ioid: ioid, sid: sid, cid: cid, updates: ch}
		c.monitors[ioid] = ms
		c.mu.Unlock()

		if err := c.send(BuildMessage(CmdMonitor, flagApp, payload.Bytes())); err != nil {
			ch <- MonitorEvent{Err: err}
		}

		// watch context cancellation → send MONITOR DESTROY
		go func() {
			<-ctx.Done()
			c.cancelMonitor(ioid, sid)
		}()
	})
}

func (c *conn) cancelMonitor(ioid, sid uint32) {
	var payload bytes.Buffer
	binary.Write(&payload, binary.LittleEndian, sid)
	binary.Write(&payload, binary.LittleEndian, ioid)
	payload.WriteByte(SubCmdDEstroy)
	_ = c.send(BuildMessage(CmdMonitor, flagApp, payload.Bytes()))
}

// ---- GET (one-shot) ---------------------------------------------------

// get issues a GET for pvName and calls cb when the response arrives.
func (c *conn) get(pvName string, cb func(v pvdata.StructValue, err error)) {
	c.openChannel(pvName, func(sid uint32, err error) {
		if err != nil {
			cb(pvdata.StructValue{}, err)
			return
		}
		ioid := c.nextIOID.Add(1)

		// INIT phase: GET with subCmd=INIT only (spec §7.2 — INIT and GET are separate messages)
		var payload bytes.Buffer
		binary.Write(&payload, binary.LittleEndian, sid)
		binary.Write(&payload, binary.LittleEndian, ioid)
		payload.WriteByte(SubCmdInit) // INIT only — server returns FieldDesc
		// pvRequest: empty struct = select all fields
		payload.WriteByte(pvdata.TypeCodeStruct)
		pvdata.WriteString(&payload, "")
		pvdata.WriteSize(&payload, 0)

		c.mu.Lock()
		ms := &monitorSub{ioid: ioid, sid: sid, updates: make(chan MonitorEvent, 1)}
		c.monitors[ioid] = ms
		c.mu.Unlock()

		if err := c.send(BuildMessage(CmdGet, flagApp, payload.Bytes())); err != nil {
			cb(pvdata.StructValue{}, err)
			return
		}

		// wait for INIT response, then issue GET
		go func() {
			evt := <-ms.updates
			if evt.Err != nil {
				cb(pvdata.StructValue{}, evt.Err)
				return
			}
			// INIT done — send GET sub-command
			var p2 bytes.Buffer
			binary.Write(&p2, binary.LittleEndian, sid)
			binary.Write(&p2, binary.LittleEndian, ioid)
			p2.WriteByte(SubCmdGet)
			if err := c.send(BuildMessage(CmdGet, flagApp, p2.Bytes())); err != nil {
				cb(pvdata.StructValue{}, err)
				return
			}
			// wait for data response
			evt = <-ms.updates
			cb(evt.Value, evt.Err)

			c.mu.Lock()
			delete(c.monitors, ioid)
			c.mu.Unlock()
		}()
	})
}

// ---- Read loop --------------------------------------------------------

func (c *conn) readLoop() {
	defer close(c.done)
	for {
		hdr, err := ReadHeader(c.br)
		if err != nil {
			if err != io.EOF {
				slog.Debug("pva readLoop", "err", err)
			}
			c.closeAllWithError(err)
			return
		}
		payload := make([]byte, hdr.Size)
		if _, err := io.ReadFull(c.br, payload); err != nil {
			c.closeAllWithError(err)
			return
		}
		r := bytes.NewReader(payload)
		bo := hdr.byteOrder()

		if hdr.isControl() {
			c.handleControl(hdr, payload)
			continue
		}
		switch hdr.Command {
		case CmdCreateChannel:
			c.handleCreateChannel(r, bo)
		case CmdGet:
			c.handleGet(r, bo)
		case CmdMonitor:
			c.handleMonitor(r, bo)
		case CmdDestroyChannel:
			// server killed the channel; clean up
		case CmdMessage:
			handleServerMessage(r, bo)
		default:
			// ignore unknown commands
		}
	}
}

func (c *conn) handleControl(hdr PVAHeader, payload []byte) {
	switch hdr.Command {
	case CtrlSetByteOrder:
		// server's byte-order announcement — we always use LE so ignore
	case CtrlEchoRequest:
		// respond with ECHO_RESPONSE
		_ = c.send(BuildMessage(CtrlEchoResponse, flagControl, payload))
	}
}

// handleCreateChannel decodes a CREATE_CHANNEL response.
// Wire: count(2) × { cid(4) + status(1+…) + sid(4) }
func (c *conn) handleCreateChannel(r io.Reader, bo binary.ByteOrder) {
	var count uint16
	binary.Read(r, bo, &count)
	for i := 0; i < int(count); i++ {
		var cid uint32
		binary.Read(r, bo, &cid)
		st, err := ReadStatus(r, bo)
		if err != nil {
			return
		}
		var sid uint32
		if st.OK() {
			binary.Read(r, bo, &sid)
		}

		c.mu.Lock()
		ch, ok := c.chans[cid]
		if !ok {
			c.mu.Unlock()
			continue
		}
		cbs := ch.onCreate
		ch.onCreate = nil
		if st.OK() {
			ch.sid = sid
			ch.state = chanCreated
		} else {
			ch.state = chanFailed
		}
		c.mu.Unlock()

		var cbErr error
		if !st.OK() {
			cbErr = fmt.Errorf("pva: %s", st.Error())
		}
		for _, cb := range cbs {
			cb(sid, cbErr)
		}
	}
}

// handleGet decodes a GET response (both INIT and data phases).
func (c *conn) handleGet(r io.Reader, bo binary.ByteOrder) {
	var ioid uint32
	binary.Read(r, bo, &ioid)
	var subCmdBuf [1]byte
	io.ReadFull(r, subCmdBuf[:])
	subCmd := subCmdBuf[0]

	st, err := ReadStatus(r, bo)
	if err != nil {
		return
	}

	c.mu.Lock()
	ms, ok := c.monitors[ioid]
	c.mu.Unlock()
	if !ok {
		return
	}

	if subCmd&SubCmdInit != 0 {
		// INIT response — read FieldDesc
		if st.OK() {
			desc, err := pvdata.ReadFieldDesc(r)
			if err == nil {
				ms.desc = desc
			}
		}
		ms.updates <- MonitorEvent{Err: asError(st)}
		return
	}

	// GET data response
	if !st.OK() {
		ms.updates <- MonitorEvent{Err: fmt.Errorf("pva GET: %s", st.Error())}
		return
	}
	// read bitSet (changed mask) then value
	changed, err := pvdata.ReadBitSet(r)
	if err != nil {
		ms.updates <- MonitorEvent{Err: err}
		return
	}
	v, err := pvdata.ReadValue(r, ms.desc)
	if err != nil {
		ms.updates <- MonitorEvent{Err: err}
		return
	}
	sv, _ := v.(pvdata.StructValue)
	ms.updates <- MonitorEvent{Changed: changed, Value: sv}
}

// handleMonitor decodes a MONITOR response.
// Sub-commands: INIT (0x08) → FieldDesc, pipeline data.
func (c *conn) handleMonitor(r io.Reader, bo binary.ByteOrder) {
	var ioid uint32
	binary.Read(r, bo, &ioid)
	var subCmdBuf2 [1]byte
	io.ReadFull(r, subCmdBuf2[:])
	subCmd := subCmdBuf2[0]

	c.mu.Lock()
	ms, ok := c.monitors[ioid]
	c.mu.Unlock()
	if !ok {
		return
	}

	if subCmd&SubCmdInit != 0 {
		st, err := ReadStatus(r, bo)
		if err != nil || !st.OK() {
			ms.updates <- MonitorEvent{Err: asError(st)}
			return
		}
		desc, err := pvdata.ReadFieldDesc(r)
		if err != nil {
			ms.updates <- MonitorEvent{Err: err}
			return
		}
		ms.desc = desc
		// Acknowledge pipeline
		c.sendMonitorAck(ioid)
		return
	}

	if subCmd&SubCmdDEstroy != 0 {
		close(ms.updates)
		c.mu.Lock()
		delete(c.monitors, ioid)
		c.mu.Unlock()
		return
	}

	// Data update: changed BitSet + overrun BitSet + value
	changed, err := pvdata.ReadBitSet(r)
	if err != nil {
		return
	}
	_, err = pvdata.ReadBitSet(r) // overrun — discard for now
	if err != nil {
		return
	}
	v, err := pvdata.ReadValue(r, ms.desc)
	if err != nil {
		return
	}
	sv, _ := v.(pvdata.StructValue)
	select {
	case ms.updates <- MonitorEvent{Changed: changed, Value: sv}:
	default:
		// slow consumer — drop (overrun)
	}
	// Acknowledge to allow more updates (pipeline)
	c.sendMonitorAck(ioid)
}

func (c *conn) sendMonitorAck(ioid uint32) {
	// MONITOR pipeline ACK: sid(4) + ioid(4) + subCmd(0x80)
	c.mu.Lock()
	ms, ok := c.monitors[ioid]
	sid := uint32(0)
	if ok {
		sid = ms.sid
	}
	c.mu.Unlock()

	var p bytes.Buffer
	binary.Write(&p, binary.LittleEndian, sid)
	binary.Write(&p, binary.LittleEndian, ioid)
	p.WriteByte(SubCmdPipeline)
	_ = c.send(BuildMessage(CmdMonitor, flagApp, p.Bytes()))
}

func handleServerMessage(r io.Reader, bo binary.ByteOrder) {
	// MESSAGE: type(1) + message(string) — log and ignore
	var t [1]byte
	r.Read(t[:])
	msg, _ := readPVAString(r, bo)
	slog.Debug("pva server message", "type", t[0], "msg", msg)
}

func (c *conn) closeAllWithError(err error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	for _, ms := range c.monitors {
		select {
		case ms.updates <- MonitorEvent{Err: err}:
		default:
		}
	}
}

func (c *conn) close() {
	c.nc.Close()
	select {
	case <-c.done:
	case <-time.After(2 * time.Second):
	}
}

func asError(s Status) error {
	if s.OK() {
		return nil
	}
	return fmt.Errorf("%s", s.Error())
}