/*
 * SPDX-License-Identifier: Apache-2.0
 * ZephCore USB CDC Companion Transport
 *
 * V3-framed USB CDC for companion mode. Extracted from main_companion.cpp.
 * Only compiled when CONFIG_LOG is enabled (debug builds).
 *
 * USBD lifecycle + 1200-baud DFU detection + DTR state tracking live in
 * the shared ZephyrUSBCDC module; this file just runs the V3 frame parser
 * on top of the CDC ACM UART and reacts to DTR-drop events from there.
 */

#include <zephyr/kernel.h>
#include <zephyr/device.h>
#include <zephyr/devicetree.h>
#include <zephyr/drivers/uart.h>
#include <zephyr/sys/ring_buffer.h>

#include <zephyr/logging/log.h>
LOG_MODULE_REGISTER(zephcore_usb, CONFIG_ZEPHCORE_USB_LOG_LEVEL);

#include <ZephyrBLE.h>
#include <app/CompanionMesh.h>

#include "ZephyrCompanionUSB.h"
#include "ZephyrUSBCDC.h"

/* MAX_FRAME_SIZE defined in CompanionMesh.h */

#define USB_RING_BUF_SIZE     512     /* USB RX ring buffer size */
#define USB_TX_RING_BUF_SIZE  2048    /* USB TX ring buffer (~13 contact frames of headroom) */
#define USB_FRAME_TIMEOUT_MS  2000    /* Partial frame timeout - reset parser after 2s of no completion */

/* Companion serial framing (MeshCore ArduinoSerialInterface):
 *   app  → device:  '<' len_lo len_hi <payload...>
 *   device → app:   '>' len_lo len_hi <payload...>
 * The leading sync byte is mandatory — the official app keys off it, and we
 * must ignore any noise/banner bytes until it arrives. */
#define USB_FRAME_RX_SYNC     '<'
#define USB_FRAME_TX_SYNC     '>'

enum usb_rx_state {
	USB_RX_IDLE = 0,  /* waiting for '<' sync byte */
	USB_RX_LEN_LO,    /* got sync, waiting len LSB */
	USB_RX_LEN_HI,    /* got len LSB, waiting len MSB */
	USB_RX_PAYLOAD,   /* accumulating payload */
};

/* USB CDC state */
static const struct device *usb_dev;
static uint8_t usb_ring_buf_data[USB_RING_BUF_SIZE];
static struct ring_buf usb_ring_buf;
static uint8_t usb_rx_buf[MAX_FRAME_SIZE];
static enum usb_rx_state usb_rx_st;
static uint16_t usb_rx_idx;     /* payload bytes received so far */
static uint16_t usb_frame_len;  /* Expected payload length (0 = none in progress) */
static uint32_t usb_frame_start_time;  /* Timestamp of sync byte for current frame */

/* TX side: interrupt-driven so the contact pump gets real backpressure +
 * a "drained" event (the USB analogue of BLE's notify-complete) instead of a
 * fixed delay. write_frame queues whole frames here under usb_tx_lock; the TX
 * ISR drains into the CDC FIFO and fires s_tx_drain_cb when the ring empties. */
static uint8_t usb_tx_ring_buf_data[USB_TX_RING_BUF_SIZE];
static struct ring_buf usb_tx_ring_buf;
static struct k_spinlock usb_tx_lock;

/* Work item for deferred V3 frame processing (set by init) */
static struct k_work *s_rx_work;

/* Session start/end callbacks (mirror BLE on_connected / on_disconnected),
 * set by main. start fires on first-frame claim, end on DTR drop. */
static void (*s_session_start_cb)(void);
static void (*s_session_end_cb)(void);

/* TX-drained callback (mirrors BLE on_tx_idle) — re-kicks the contact pump. */
static void (*s_tx_drain_cb)(void);

/* Work items */
static void usb_rx_work_fn(struct k_work *work);

K_WORK_DEFINE(usb_rx_work, usb_rx_work_fn);

/* USB CDC UART interrupt callback - puts bytes in ring buffer */
static void usb_uart_isr(const struct device *dev, void *user_data)
{
	ARG_UNUSED(user_data);

	while (uart_irq_update(dev) && uart_irq_is_pending(dev)) {
		if (uart_irq_rx_ready(dev)) {
			uint8_t buf[64];
			int recv_len = uart_fifo_read(dev, buf, sizeof(buf));
			if (recv_len > 0) {
				ring_buf_put(&usb_ring_buf, buf, recv_len);
				k_work_submit(&usb_rx_work);
			}
		}

		if (uart_irq_tx_ready(dev)) {
			/* Push as much of our TX ring as the CDC FIFO will take.
			 * uart_fifo_fill returns the count actually accepted, so the
			 * remainder stays queued and the callback fires again when the
			 * FIFO drains — backpressure all the way to the wire. */
			uint8_t *out;
			bool empty;
			k_spinlock_key_t key = k_spin_lock(&usb_tx_lock);
			uint32_t claimed = ring_buf_get_claim(&usb_tx_ring_buf, &out, 64);
			if (claimed > 0) {
				int sent = uart_fifo_fill(dev, out, claimed);
				ring_buf_get_finish(&usb_tx_ring_buf, sent > 0 ? sent : 0);
			}
			empty = ring_buf_is_empty(&usb_tx_ring_buf);
			if (empty) {
				/* Disable inside the lock so a concurrent write_frame can't
				 * enqueue+enable in the gap and then have us disable it,
				 * stranding the frame. write_frame's enable always runs
				 * after its put, so it re-arms us correctly. */
				uart_irq_tx_disable(dev);
			}
			k_spin_unlock(&usb_tx_lock, key);

			if (empty && s_tx_drain_cb) {
				/* Channel idle — let the pump queue the next batch. */
				s_tx_drain_cb();
			}
		}
	}
}

/* USB RX work - parses V3 frames from ring buffer */
static void usb_rx_work_fn(struct k_work *work)
{
	ARG_UNUSED(work);

	uint8_t byte;

	/* Timeout partial frames — if we've been mid-frame too long without
	 * completing, reset the parser state and resync on the next sync byte. */
	if (usb_rx_st != USB_RX_IDLE &&
	    (k_uptime_get_32() - usb_frame_start_time) > USB_FRAME_TIMEOUT_MS) {
		LOG_WRN("usb_rx: partial frame timeout (state=%d, expected=%u), resync",
			usb_rx_st, usb_frame_len);
		usb_rx_st = USB_RX_IDLE;
		usb_frame_len = 0;
		usb_rx_idx = 0;
	}

	while (ring_buf_get(&usb_ring_buf, &byte, 1) == 1) {
		switch (usb_rx_st) {
		case USB_RX_IDLE:
			/* Ignore stray bytes until the '<' sync byte arrives. */
			if (byte == USB_FRAME_RX_SYNC) {
				usb_rx_st = USB_RX_LEN_LO;
				usb_frame_start_time = k_uptime_get_32();
			}
			break;
		case USB_RX_LEN_LO:
			usb_frame_len = byte;  /* LSB */
			usb_rx_st = USB_RX_LEN_HI;
			break;
		case USB_RX_LEN_HI:
			usb_frame_len |= ((uint16_t)byte) << 8;  /* MSB */
			usb_rx_idx = 0;
			if (usb_frame_len == 0 || usb_frame_len > MAX_FRAME_SIZE) {
				LOG_WRN("usb_rx: invalid frame len %u, resync", usb_frame_len);
				usb_rx_st = USB_RX_IDLE;
				usb_frame_len = 0;
			} else {
				usb_rx_st = USB_RX_PAYLOAD;
			}
			break;
		default: /* USB_RX_PAYLOAD */
			usb_rx_buf[usb_rx_idx++] = byte;

			if (usb_rx_idx >= usb_frame_len) {
				/* Frame complete - queue it */
				uint8_t *payload = usb_rx_buf;
				uint16_t payload_len = usb_frame_len;

				LOG_DBG("usb_rx: frame complete len=%u hdr=0x%02x", payload_len, payload[0]);

				/* Claim the interface for USB on the FIRST inbound frame of
				 * any opcode — mirroring BLE, which claims on connect. The
				 * official client opens with CMD_DEVICE_QUERY (0x16), not
				 * CMD_APP_START (0x01); gating the claim on 0x01 dropped that
				 * first query and the client timed out waiting for a reply. */
				if (zephcore_ble_get_active_iface() != ZEPHCORE_IFACE_USB) {
					/* Atomically claim the interface for USB unless BLE
					 * already owns it.  try_claim succeeds when idle or
					 * already USB (reconnect) and fails only while a BLE
					 * session is live, so USB can't steal it — and the
					 * compare-and-set can't race a concurrent BLE claim. */
					if (zephcore_ble_iface_try_claim(ZEPHCORE_IFACE_USB)) {
						zephcore_ble_set_enabled(false);
						LOG_INF("usb_rx: first frame 0x%02x → IFACE_USB", payload[0]);
						/* New USB session — mirror BLE's on-connect UI
						 * notification (Arduino shows "connected" for serial
						 * transports too). Fires once per session: try_claim
						 * only returns true on the NONE→USB transition. */
						if (s_session_start_cb) {
							s_session_start_cb();
						}
					} else {
						LOG_INF("usb_rx: frame 0x%02x ignored, BLE is active", payload[0]);
					}
				}

				/* Only process if USB is active interface */
				if (zephcore_ble_get_active_iface() == ZEPHCORE_IFACE_USB) {
					struct {
						uint16_t len;
						uint8_t buf[MAX_FRAME_SIZE];
					} f;
					f.len = payload_len;
					memcpy(f.buf, payload, payload_len);
					/* sysworkq handles V3-protocol parsing; main only wakes
					 * if downstream LoRa work gets enqueued (via TX_DRAIN). */
					if (k_msgq_put(zephcore_ble_get_recv_queue(), &f, K_NO_WAIT) == 0) {
						k_work_submit(s_rx_work);
					}
				}

				/* Reset for next frame */
				usb_rx_st = USB_RX_IDLE;
				usb_frame_len = 0;
				usb_rx_idx = 0;
			}
			break;
		}
	}
}

/* DTR-transition callback from the shared ZephyrUSBCDC module.
 * On drop: host closed the port → reset parser, hand control back to BLE. */
static void on_dtr_change(bool dtr_active)
{
	if (dtr_active) {
		return;
	}
	LOG_INF("usb_dtr: DTR dropped, USB disconnected");
	/* While USB owns the interface, BLE claims are rejected (see connected()),
	 * so this thread is the only writer of active_iface here — the get/set
	 * pair below needs no extra locking beyond the thread-safe accessors. */
	if (zephcore_ble_get_active_iface() == ZEPHCORE_IFACE_USB) {
		/* A USB companion session is ending — run the same per-session
		 * cleanup BLE does on disconnect (cancel contact iteration and
		 * message sync, free the sign buffer). Without this, stale sync
		 * state carries into the next session and an in-flight sign op
		 * leaks its 8KB buffer. */
		if (s_session_end_cb) {
			s_session_end_cb();
		}
		if (zephcore_ble_is_connected()) {
			/* BLE client is physically connected — hand off to it. */
			zephcore_ble_set_active_iface(ZEPHCORE_IFACE_BLE);
			LOG_INF("usb_dtr: → IFACE_BLE (BLE was connected)");
		} else {
			/* Nobody connected — go idle and restart advertising
			 * so the phone can find the companion again. */
			zephcore_ble_set_active_iface(ZEPHCORE_IFACE_NONE);
			zephcore_ble_set_enabled(true);
			LOG_INF("usb_dtr: → IFACE_NONE, BLE advertising restarted");
		}
	}
	ring_buf_reset(&usb_ring_buf);
	usb_rx_st = USB_RX_IDLE;
	usb_frame_len = 0;
	usb_rx_idx = 0;

	/* Discard any pending TX from the closed session. */
	uart_irq_tx_disable(usb_dev);
	k_spinlock_key_t key = k_spin_lock(&usb_tx_lock);
	ring_buf_reset(&usb_tx_ring_buf);
	k_spin_unlock(&usb_tx_lock, key);
}

/* Queue a frame for interrupt-driven TX (sync byte + LE length + payload,
 * matching the MeshCore ArduinoSerialInterface framing). The whole frame is
 * committed atomically under usb_tx_lock — either it all fits or none of it
 * does (returns 0), so frames never tear and tx_has_space() stays truthful.
 * 0 means "ring full, retry when drained"; the caller (contact pump) backs off
 * and the TX-drain callback re-kicks it. */
size_t zephcore_usb_companion_write_frame(const uint8_t *src, size_t len)
{
	if (!usb_dev || len == 0 || len > MAX_FRAME_SIZE) {
		return 0;
	}

	uint8_t hdr[3] = {
		USB_FRAME_TX_SYNC,
		(uint8_t)(len & 0xFF),
		(uint8_t)((len >> 8) & 0xFF),
	};
	size_t total = sizeof(hdr) + len;

	k_spinlock_key_t key = k_spin_lock(&usb_tx_lock);
	if (ring_buf_space_get(&usb_tx_ring_buf) < total) {
		k_spin_unlock(&usb_tx_lock, key);
		return 0;
	}
	ring_buf_put(&usb_tx_ring_buf, hdr, sizeof(hdr));
	ring_buf_put(&usb_tx_ring_buf, src, len);
	k_spin_unlock(&usb_tx_lock, key);

	/* Kick the TX ISR; harmless if already enabled. */
	uart_irq_tx_enable(usb_dev);

	LOG_DBG("usb_write_frame: queued len=%u hdr=0x%02x", (unsigned)len, src[0]);
	return len;
}

/* True if the TX ring can hold one more frame of `payload_len` (+3 framing).
 * The pump checks this before each contact so write_frame can't fail mid-dump. */
bool zephcore_usb_companion_tx_has_space(size_t payload_len)
{
	if (!usb_dev) {
		return false;
	}
	k_spinlock_key_t key = k_spin_lock(&usb_tx_lock);
	bool ok = ring_buf_space_get(&usb_tx_ring_buf) >= payload_len + 3;
	k_spin_unlock(&usb_tx_lock, key);
	return ok;
}

void zephcore_usb_companion_reset_rx(void)
{
	ring_buf_reset(&usb_ring_buf);
	usb_rx_st = USB_RX_IDLE;
	usb_frame_len = 0;
	usb_rx_idx = 0;

	/* Drop any half-sent TX too — the session it belonged to is gone. */
	if (usb_dev) {
		uart_irq_tx_disable(usb_dev);
	}
	k_spinlock_key_t key = k_spin_lock(&usb_tx_lock);
	ring_buf_reset(&usb_tx_ring_buf);
	k_spin_unlock(&usb_tx_lock, key);
}

void zephcore_usb_companion_set_session_start_cb(void (*cb)(void))
{
	s_session_start_cb = cb;
}

void zephcore_usb_companion_set_session_end_cb(void (*cb)(void))
{
	s_session_end_cb = cb;
}

void zephcore_usb_companion_set_tx_drain_cb(void (*cb)(void))
{
	s_tx_drain_cb = cb;
}

void zephcore_usb_companion_init(struct k_event *mesh_events,
				 struct k_work *rx_work,
				 uint32_t mesh_event_ble_rx,
				 void *board)
{
	ARG_UNUSED(board);
	ARG_UNUSED(mesh_events);
	ARG_UNUSED(mesh_event_ble_rx);

	s_rx_work = rx_work;

	/* The cdc_acm_uart DT node may be present without the class driver compiled
	 * (shared esp32s3_usb_otg.dtsi exposes the node unconditionally; the class is
	 * only enabled with esp32s3_usb.conf). Gate the device-get on the class too,
	 * else DEVICE_DT_GET_ONE references an undefined device ordinal on, e.g., a
	 * debug ESP32-S3 companion (CONFIG_LOG=y) built without esp32s3_usb.conf. */
#if DT_HAS_COMPAT_STATUS_OKAY(zephyr_cdc_acm_uart) && \
	(IS_ENABLED(CONFIG_USB_CDC_ACM) || IS_ENABLED(CONFIG_USBD_CDC_ACM_CLASS))
	usb_dev = DEVICE_DT_GET_ONE(zephyr_cdc_acm_uart);
	if (device_is_ready(usb_dev)) {
		LOG_INF("USB CDC device ready: %s", usb_dev->name);
		ring_buf_init(&usb_ring_buf, sizeof(usb_ring_buf_data), usb_ring_buf_data);
		ring_buf_init(&usb_tx_ring_buf, sizeof(usb_tx_ring_buf_data), usb_tx_ring_buf_data);

		/* Set up UART interrupt callback (RX enabled now, TX enabled on demand
		 * by write_frame and disabled by the ISR when the TX ring drains). */
		uart_irq_callback_set(usb_dev, usb_uart_isr);
		uart_irq_rx_enable(usb_dev);

		/* DTR state changes (including disconnect) reach us via the
		 * shared usbd_msg_callback — no polling work needed. */
		zephcore_usbd_set_dtr_cb(on_dtr_change);
	} else {
		LOG_WRN("USB CDC device not ready");
		usb_dev = NULL;
	}
#endif
}