use crate::config::Config;
use crate::proto;
use anyhow::{Context, Result};
use prost::Message;
use tracing::{debug, error, info, warn};

const PROTOCOL_VERSION: u8 = 0x01;
const MSG_TYPE_SUBMIT_REQUEST: u8 = 0x10;
const MSG_TYPE_SUBMIT_RESPONSE: u8 = 0x11;
const MSG_TYPE_DATA_REQUEST: u8 = 0x01;
const MSG_TYPE_HISTORY_READY: u8 = 0x12;

pub struct Relay {
    config: Config,
    context: zmq::Context,
}

impl Relay {
    pub fn new(config: Config) -> Result<Self> {
        let context = zmq::Context::new();

        Ok(Self {
            config,
            context,
        })
    }

    pub async fn run(self) -> Result<()> {
        info!("Initializing Stateless ZMQ Relay");

        // Bind sockets
        let client_request_socket = self.create_client_request_socket()?;
        let market_data_frontend = self.create_market_data_frontend()?;
        let market_data_backend = self.create_market_data_backend()?;
        let ingestor_work_socket = self.create_ingestor_work_socket()?;

        info!("All sockets initialized successfully - relay is STATELESS");
        info!("No pending requests tracked - all async via pub/sub");

        // Run main loop
        tokio::task::spawn_blocking(move || {
            Self::proxy_loop(
                client_request_socket,
                market_data_frontend,
                market_data_backend,
                ingestor_work_socket,
            )
        })
        .await?
    }

    fn create_client_request_socket(&self) -> Result<zmq::Socket> {
        let socket = self.context.socket(zmq::ROUTER)?;
        socket.set_sndhwm(self.config.high_water_mark)?;
        socket.set_rcvhwm(self.config.high_water_mark)?;
        socket.set_linger(1000)?;

        let endpoint = format!("{}:{}", self.config.bind_address, self.config.client_request_port);
        socket.bind(&endpoint)?;
        info!("Client request socket (ROUTER) bound to {}", endpoint);
        info!("  → Accepts SubmitHistoricalRequest, returns SubmitResponse immediately");

        Ok(socket)
    }

    fn create_market_data_frontend(&self) -> Result<zmq::Socket> {
        let socket = self.context.socket(zmq::XPUB)?;
        socket.set_sndhwm(self.config.high_water_mark)?;
        socket.set_xpub_verbose(true)?;

        let endpoint = format!("{}:{}", self.config.bind_address, self.config.market_data_pub_port);
        socket.bind(&endpoint)?;
        info!("Market data frontend (XPUB) bound to {}", endpoint);
        info!("  → Clients subscribe here for HistoryReadyNotification and market data");

        Ok(socket)
    }

    fn create_market_data_backend(&self) -> Result<zmq::Socket> {
        let socket = self.context.socket(zmq::XSUB)?;
        socket.set_rcvhwm(self.config.high_water_mark)?;

        socket.connect(&self.config.flink_market_data_endpoint)?;
        info!("Market data backend (XSUB) connected to {}", self.config.flink_market_data_endpoint);
        info!("  → Receives HistoryReadyNotification and market data from Flink");

        Ok(socket)
    }

    fn create_ingestor_work_socket(&self) -> Result<zmq::Socket> {
        let socket = self.context.socket(zmq::PUB)?;
        socket.set_sndhwm(self.config.high_water_mark)?;
        socket.set_linger(1000)?;

        let endpoint = format!("{}:{}", self.config.bind_address, self.config.ingestor_work_port);
        socket.bind(&endpoint)?;
        info!("Ingestor work queue (PUB) bound to {}", endpoint);
        info!("  → Publishes DataRequest with exchange prefix");

        Ok(socket)
    }

    fn proxy_loop(
        client_request_socket: zmq::Socket,
        market_data_frontend: zmq::Socket,
        market_data_backend: zmq::Socket,
        ingestor_work_socket: zmq::Socket,
    ) -> Result<()> {
        let mut items = [
            client_request_socket.as_poll_item(zmq::POLLIN),
            market_data_frontend.as_poll_item(zmq::POLLIN),
            market_data_backend.as_poll_item(zmq::POLLIN),
        ];

        info!("Entering stateless proxy loop");

        loop {
            // Poll with 100ms timeout
            zmq::poll(&mut items, 100)
                .context("Failed to poll sockets")?;

            // Handle client request submissions
            if items[0].is_readable() {
                if let Err(e) = Self::handle_client_submission(
                    &client_request_socket,
                    &ingestor_work_socket,
                ) {
                    error!("Error handling client submission: {}", e);
                }
            }

            // Handle market data subscriptions from clients (XPUB → XSUB)
            if items[1].is_readable() {
                if let Err(e) = Self::proxy_subscription(&market_data_frontend, &market_data_backend) {
                    error!("Error proxying subscription: {}", e);
                }
            }

            // Handle market data from Flink (XSUB → XPUB)
            // This includes HistoryReadyNotification and regular market data
            if items[2].is_readable() {
                if let Err(e) = Self::proxy_market_data(&market_data_backend, &market_data_frontend) {
                    error!("Error proxying market data: {}", e);
                }
            }
        }
    }

    fn handle_client_submission(
        client_socket: &zmq::Socket,
        ingestor_socket: &zmq::Socket,
    ) -> Result<()> {
        // Receive from client: [identity][empty][version][message]
        let identity = client_socket.recv_bytes(0)?;
        let _empty = client_socket.recv_bytes(0)?;
        let version_frame = client_socket.recv_bytes(0)?;
        let message_frame = client_socket.recv_bytes(0)?;

        if version_frame.len() != 1 || version_frame[0] != PROTOCOL_VERSION {
            warn!("Invalid protocol version from client");
            return Ok(());
        }

        if message_frame.is_empty() {
            warn!("Empty message frame from client");
            return Ok(());
        }

        let msg_type = message_frame[0];
        let payload = &message_frame[1..];

        debug!("Received client submission: type=0x{:02x}, payload_len={}", msg_type, payload.len());

        match msg_type {
            MSG_TYPE_SUBMIT_REQUEST => {
                Self::handle_submit_request(
                    identity,
                    payload,
                    client_socket,
                    ingestor_socket,
                )?;
            }
            _ => {
                warn!("Unknown message type from client: 0x{:02x}", msg_type);
            }
        }

        Ok(())
    }

    fn handle_submit_request(
        client_identity: Vec<u8>,
        payload: &[u8],
        client_socket: &zmq::Socket,
        ingestor_socket: &zmq::Socket,
    ) -> Result<()> {
        // Parse protobuf request
        let request = proto::SubmitHistoricalRequest::decode(payload)
            .context("Failed to parse SubmitHistoricalRequest")?;

        let request_id = request.request_id.clone();
        let ticker = request.ticker.clone();
        let client_id = request.client_id.clone();

        info!("Handling request submission: request_id={}, ticker={}, client_id={:?}",
            request_id, ticker, client_id);

        // Extract exchange prefix from ticker
        let exchange_prefix = ticker.split(':').next()
            .map(|s| format!("{}:", s))
            .unwrap_or_else(|| String::from(""));

        if exchange_prefix.is_empty() {
            warn!("Ticker '{}' missing exchange prefix", ticker);
        }

        // Build DataRequest protobuf for ingestors
        let data_request = proto::DataRequest {
            request_id: request_id.clone(),
            r#type: proto::data_request::RequestType::HistoricalOhlc as i32,
            ticker: ticker.clone(),
            historical: Some(proto::HistoricalParams {
                start_time: request.start_time,
                end_time: request.end_time,
                period_seconds: request.period_seconds,
                limit: request.limit,
            }),
            realtime: None,
            client_id: client_id.clone(),
        };

        let mut data_request_bytes = Vec::new();
        data_request.encode(&mut data_request_bytes)?;

        // Publish to ingestors with exchange prefix
        let version_frame = vec![PROTOCOL_VERSION];
        let mut message_frame = vec![MSG_TYPE_DATA_REQUEST];
        message_frame.extend_from_slice(&data_request_bytes);

        ingestor_socket.send(&exchange_prefix, zmq::SNDMORE)?;
        ingestor_socket.send(&version_frame, zmq::SNDMORE)?;
        ingestor_socket.send(&message_frame, 0)?;

        info!("Published to ingestors: prefix={}, request_id={}", exchange_prefix, request_id);

        // Build SubmitResponse protobuf
        // NOTE: This topic is DETERMINISTIC based on client-generated values.
        // Client should have already subscribed to this topic BEFORE sending the request
        // to prevent race condition where notification arrives before client subscribes.
        let notification_topic = if let Some(cid) = &client_id {
            format!("RESPONSE:{}", cid)
        } else {
            format!("HISTORY_READY:{}", request_id)
        };

        let response = proto::SubmitResponse {
            request_id: request_id.clone(),
            status: proto::submit_response::SubmitStatus::Queued as i32,
            error_message: None,
            notification_topic: notification_topic.clone(),
        };

        let mut response_bytes = Vec::new();
        response.encode(&mut response_bytes)?;

        // Send immediate response to client
        let version_frame = vec![PROTOCOL_VERSION];
        let mut message_frame = vec![MSG_TYPE_SUBMIT_RESPONSE];
        message_frame.extend_from_slice(&response_bytes);

        client_socket.send(&client_identity, zmq::SNDMORE)?;
        client_socket.send(&[] as &[u8], zmq::SNDMORE)?;
        client_socket.send(&version_frame, zmq::SNDMORE)?;
        client_socket.send(&message_frame, 0)?;

        info!("Sent SubmitResponse to client: request_id={}, topic={}", request_id, notification_topic);

        // Relay is now DONE with this request - completely stateless!
        // Client will receive notification via pub/sub when Flink publishes HistoryReadyNotification

        Ok(())
    }

    fn proxy_subscription(
        frontend: &zmq::Socket,
        backend: &zmq::Socket,
    ) -> Result<()> {
        // Forward subscription message from XPUB to XSUB
        let msg = frontend.recv_bytes(0)?;
        backend.send(&msg, 0)?;

        if !msg.is_empty() {
            let action = if msg[0] == 1 { "subscribe" } else { "unsubscribe" };
            let topic = String::from_utf8_lossy(&msg[1..]);
            debug!("Client {} to topic: {}", action, topic);
        }

        Ok(())
    }

    fn proxy_market_data(
        backend: &zmq::Socket,
        frontend: &zmq::Socket,
    ) -> Result<()> {
        // Forward all messages from XSUB to XPUB (zero-copy proxy)
        // This includes:
        // - Regular market data (ticks, OHLC)
        // - HistoryReadyNotification from Flink
        loop {
            let msg = backend.recv_bytes(0)?;
            let more = backend.get_rcvmore()?;

            if more {
                frontend.send(&msg, zmq::SNDMORE)?;
            } else {
                frontend.send(&msg, 0)?;
                break;
            }
        }

        Ok(())
    }
}
