use std::{
    collections::HashMap,
    sync::{
        atomic::{AtomicBool, AtomicU64, Ordering},
        Arc,
    },
    time::Duration,
};

use alloy::primitives::Address;
use futures_util::{stream::SplitSink, SinkExt, StreamExt};
use log::{error, info, warn};
use serde::{Deserialize, Serialize};
use tokio::{
    net::TcpStream,
    spawn,
    sync::{mpsc, mpsc::UnboundedSender, oneshot, Mutex},
    time,
};
use tokio_tungstenite::{
    connect_async,
    tungstenite::{self, protocol},
    MaybeTlsStream, WebSocketStream,
};

use crate::{
    prelude::*,
    ws::message_types::{
        ActiveAssetData, ActiveSpotAssetCtx, AllMids, Bbo, Candle, L2Book, OrderUpdates, Trades,
        User,
    },
    ActiveAssetCtx, Error, ExchangeApiError, ExchangeResponseStatus, Notification, UserFills,
    UserFundings, UserNonFundingLedgerUpdates, WebData2,
};

type WsWriter = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, protocol::Message>;

#[derive(Debug)]
struct SubscriptionData {
    sending_channel: UnboundedSender<Message>,
    subscription_id: u32,
    id: String,
}
#[derive(Debug)]
pub(crate) struct WsManager {
    stop_flag: Arc<AtomicBool>,
    write_tx: mpsc::UnboundedSender<protocol::Message>,
    subscriptions: Arc<Mutex<HashMap<String, Vec<SubscriptionData>>>>,
    subscription_id: u32,
    subscription_identifiers: HashMap<u32, String>,
    pending_posts: Arc<Mutex<HashMap<u64, oneshot::Sender<ExchangeResponseStatus>>>>,
    next_post_id: AtomicU64,
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "type")]
#[serde(rename_all = "camelCase")]
pub enum Subscription {
    AllMids,
    Notification { user: Address },
    WebData2 { user: Address },
    Candle { coin: String, interval: String },
    L2Book { coin: String },
    Trades { coin: String },
    OrderUpdates { user: Address },
    UserEvents { user: Address },
    UserFills { user: Address },
    UserFundings { user: Address },
    UserNonFundingLedgerUpdates { user: Address },
    ActiveAssetCtx { coin: String },
    ActiveAssetData { user: Address, coin: String },
    Bbo { coin: String },
}

#[derive(Deserialize, Clone, Debug)]
#[serde(tag = "channel")]
#[serde(rename_all = "camelCase")]
pub enum Message {
    NoData,
    HyperliquidError(String),
    AllMids(AllMids),
    Trades(Trades),
    L2Book(L2Book),
    User(User),
    UserFills(UserFills),
    Candle(Candle),
    SubscriptionResponse,
    OrderUpdates(OrderUpdates),
    UserFundings(UserFundings),
    UserNonFundingLedgerUpdates(UserNonFundingLedgerUpdates),
    Notification(Notification),
    WebData2(WebData2),
    ActiveAssetCtx(ActiveAssetCtx),
    ActiveAssetData(ActiveAssetData),
    ActiveSpotAssetCtx(ActiveSpotAssetCtx),
    Bbo(Bbo),
    Post(PostResponseEnvelope),
    Pong,
}

/// Envelope for WS post responses: `{"channel": "post", "data": {"id": N, "response": {...}}}`
#[derive(Deserialize, Clone, Debug)]
pub struct PostResponseEnvelope {
    pub data: PostResponseData,
}

#[derive(Deserialize, Clone, Debug)]
pub struct PostResponseData {
    pub id: u64,
    pub response: PostResponsePayload,
}

#[derive(Deserialize, Clone, Debug)]
#[serde(tag = "type")]
#[serde(rename_all = "camelCase")]
pub enum PostResponsePayload {
    Action { payload: ExchangeResponseStatus },
    Info { payload: serde_json::Value },
    Error(PostError),
}

#[derive(Deserialize, Clone, Debug)]
pub struct PostError {
    pub payload: String,
}

#[derive(Serialize)]
pub(crate) struct SubscriptionSendData<'a> {
    method: &'static str,
    subscription: &'a serde_json::Value,
}

#[derive(Serialize)]
pub(crate) struct Ping {
    method: &'static str,
}

/// Envelope for sending a post request over WS.
#[derive(Serialize, Debug)]
pub struct WsPostSendData {
    pub method: &'static str,
    pub id: u64,
    pub request: WsPostRequest,
}

#[derive(Serialize, Debug)]
pub struct WsPostRequest {
    #[serde(rename = "type")]
    pub request_type: &'static str,
    pub payload: serde_json::Value,
}

impl WsManager {
    const SEND_PING_INTERVAL: u64 = 50;

    pub(crate) async fn new(url: String, reconnect: bool) -> Result<WsManager> {
        let stop_flag = Arc::new(AtomicBool::new(false));

        let (writer, reader) = Self::connect(&url).await?.split();

        let subscriptions: Arc<Mutex<HashMap<String, Vec<SubscriptionData>>>> =
            Arc::new(Mutex::new(HashMap::new()));

        let pending_posts: Arc<Mutex<HashMap<u64, oneshot::Sender<ExchangeResponseStatus>>>> =
            Arc::new(Mutex::new(HashMap::new()));

        let (write_tx, write_rx) = mpsc::unbounded_channel::<protocol::Message>();
        let writer = Arc::new(Mutex::new(writer));

        Self::spawn_writer_task(Arc::clone(&writer), write_rx, Arc::clone(&stop_flag));
        Self::spawn_reader_task(
            url,
            reconnect,
            reader,
            Arc::clone(&writer),
            Arc::clone(&subscriptions),
            Arc::clone(&pending_posts),
            Arc::clone(&stop_flag),
        );
        Self::spawn_ping_task(write_tx.clone(), Arc::clone(&stop_flag));

        Ok(WsManager {
            stop_flag,
            write_tx,
            subscriptions,
            subscription_id: 0,
            subscription_identifiers: HashMap::new(),
            pending_posts,
            next_post_id: AtomicU64::new(1),
        })
    }

    fn spawn_writer_task(
        writer: Arc<Mutex<WsWriter>>,
        mut write_rx: mpsc::UnboundedReceiver<protocol::Message>,
        stop_flag: Arc<AtomicBool>,
    ) {
        spawn(async move {
            while !stop_flag.load(Ordering::Relaxed) {
                match write_rx.recv().await {
                    Some(msg) => {
                        let mut w = writer.lock().await;
                        if let Err(err) = w.send(msg).await {
                            error!("WsManager writer error: {err}");
                        }
                    }
                    None => break,
                }
            }
            warn!("ws writer task stopped");
        });
    }

    fn spawn_reader_task(
        url: String,
        reconnect: bool,
        mut reader: futures_util::stream::SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
        writer: Arc<Mutex<WsWriter>>,
        subscriptions: Arc<Mutex<HashMap<String, Vec<SubscriptionData>>>>,
        pending_posts: Arc<Mutex<HashMap<u64, oneshot::Sender<ExchangeResponseStatus>>>>,
        stop_flag: Arc<AtomicBool>,
    ) {
        spawn(async move {
            while !stop_flag.load(Ordering::Relaxed) {
                if let Some(data) = reader.next().await {
                    if let Err(err) =
                        WsManager::parse_and_send_data(data, &subscriptions, &pending_posts).await
                    {
                        error!("Error processing data received by WsManager reader: {err}");
                    }
                } else {
                    warn!("WsManager disconnected");
                    {
                        let mut pending = pending_posts.lock().await;
                        for (_, sender) in pending.drain() {
                            drop(sender);
                        }
                    }
                    if let Err(err) =
                        WsManager::send_to_all_subscriptions(&subscriptions, Message::NoData).await
                    {
                        warn!("Error sending disconnection notification err={err}");
                    }
                    if reconnect {
                        Self::handle_reconnect(&url, &mut reader, &writer, &subscriptions).await;
                    } else {
                        error!("WsManager reconnection disabled. Will not reconnect and exiting reader task.");
                        break;
                    }
                }
            }
            warn!("ws message reader task stopped");
        });
    }

    async fn handle_reconnect(
        url: &str,
        reader: &mut futures_util::stream::SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
        writer: &Arc<Mutex<WsWriter>>,
        subscriptions: &Arc<Mutex<HashMap<String, Vec<SubscriptionData>>>>,
    ) {
        tokio::time::sleep(Duration::from_secs(1)).await;
        info!("WsManager attempting to reconnect");
        match Self::connect(url).await {
            Ok(ws) => {
                let (new_writer, new_reader) = ws.split();
                *reader = new_reader;
                let mut writer_guard = writer.lock().await;
                *writer_guard = new_writer;
                for (identifier, v) in subscriptions.lock().await.iter() {
                    if identifier.eq("userEvents") || identifier.eq("orderUpdates") {
                        for subscription_data in v {
                            if let Err(err) = Self::send_subscription_data_to_writer(
                                "subscribe",
                                &mut writer_guard,
                                &subscription_data.id,
                            )
                            .await
                            {
                                error!("Could not resubscribe {identifier}: {err}");
                            }
                        }
                    } else if let Err(err) = Self::send_subscription_data_to_writer(
                        "subscribe",
                        &mut writer_guard,
                        identifier,
                    )
                    .await
                    {
                        error!("Could not resubscribe correctly {identifier}: {err}");
                    }
                }
                info!("WsManager reconnect finished");
            }
            Err(err) => error!("Could not connect to websocket {err}"),
        }
    }

    fn spawn_ping_task(
        write_tx: mpsc::UnboundedSender<protocol::Message>,
        stop_flag: Arc<AtomicBool>,
    ) {
        spawn(async move {
            while !stop_flag.load(Ordering::Relaxed) {
                match serde_json::to_string(&Ping { method: "ping" }) {
                    Ok(payload) => {
                        if let Err(err) = write_tx.send(protocol::Message::Text(payload)) {
                            error!("Error sending ping to writer channel: {err}")
                        }
                    }
                    Err(err) => error!("Error serializing ping message: {err}"),
                }
                time::sleep(Duration::from_secs(Self::SEND_PING_INTERVAL)).await;
            }
            warn!("ws ping task stopped");
        });
    }

    async fn connect(url: &str) -> Result<WebSocketStream<MaybeTlsStream<TcpStream>>> {
        let (ws_stream, _) = connect_async(url)
            .await
            .map_err(|e| Error::Websocket(e.to_string()))?;

        // Disable Nagle's algorithm to send small messages (cancels, orders)
        // immediately rather than buffering them.
        let tcp_stream = ws_stream.get_ref();
        match tcp_stream {
            MaybeTlsStream::Plain(s) => {
                s.set_nodelay(true).ok();
            }
            MaybeTlsStream::NativeTls(s) => {
                let tcp: &TcpStream = s.get_ref().get_ref().get_ref();
                tcp.set_nodelay(true).ok();
            }
            _ => {}
        }

        Ok(ws_stream)
    }

    pub(crate) async fn post_action(
        &self,
        payload: serde_json::Value,
    ) -> Result<ExchangeResponseStatus> {
        let id = self.next_post_id.fetch_add(1, Ordering::Relaxed);

        let send_data = WsPostSendData {
            method: "post",
            id,
            request: WsPostRequest {
                request_type: "action",
                payload,
            },
        };

        let payload_str =
            serde_json::to_string(&send_data).map_err(|e| Error::JsonParse(e.to_string()))?;

        let (tx, rx) = oneshot::channel();
        self.pending_posts.lock().await.insert(id, tx);

        if let Err(e) = self.write_tx.send(protocol::Message::Text(payload_str)) {
            self.pending_posts.lock().await.remove(&id);
            return Err(Error::WsPostFailed(format!("writer channel closed: {e}")));
        }

        match time::timeout(Duration::from_secs(5), rx).await {
            Ok(Ok(response)) => Ok(response),
            Ok(Err(_)) => {
                // oneshot sender dropped (WS disconnected and drained pending)
                Err(Error::WsPostResponseDropped)
            }
            Err(_) => {
                // Timeout — clean up the pending entry
                self.pending_posts.lock().await.remove(&id);
                Err(Error::WsPostFailed(
                    "post request timed out after 5s".to_string(),
                ))
            }
        }
    }

    fn get_identifier(message: &Message) -> Result<String> {
        match message {
            Message::AllMids(_) => serde_json::to_string(&Subscription::AllMids)
                .map_err(|e| Error::JsonParse(e.to_string())),
            Message::User(_) => Ok("userEvents".to_string()),
            Message::UserFills(fills) => serde_json::to_string(&Subscription::UserFills {
                user: fills.data.user,
            })
            .map_err(|e| Error::JsonParse(e.to_string())),
            Message::Trades(trades) => {
                if trades.data.is_empty() {
                    Ok(String::default())
                } else {
                    serde_json::to_string(&Subscription::Trades {
                        coin: trades.data[0].coin.clone(),
                    })
                    .map_err(|e| Error::JsonParse(e.to_string()))
                }
            }
            Message::L2Book(l2_book) => serde_json::to_string(&Subscription::L2Book {
                coin: l2_book.data.coin.clone(),
            })
            .map_err(|e| Error::JsonParse(e.to_string())),
            Message::Candle(candle) => serde_json::to_string(&Subscription::Candle {
                coin: candle.data.coin.clone(),
                interval: candle.data.interval.clone(),
            })
            .map_err(|e| Error::JsonParse(e.to_string())),
            Message::OrderUpdates(_) => Ok("orderUpdates".to_string()),
            Message::UserFundings(fundings) => serde_json::to_string(&Subscription::UserFundings {
                user: fundings.data.user,
            })
            .map_err(|e| Error::JsonParse(e.to_string())),
            Message::UserNonFundingLedgerUpdates(user_non_funding_ledger_updates) => {
                serde_json::to_string(&Subscription::UserNonFundingLedgerUpdates {
                    user: user_non_funding_ledger_updates.data.user,
                })
                .map_err(|e| Error::JsonParse(e.to_string()))
            }
            Message::Notification(_) => Ok("notification".to_string()),
            Message::WebData2(web_data2) => serde_json::to_string(&Subscription::WebData2 {
                user: web_data2.data.user,
            })
            .map_err(|e| Error::JsonParse(e.to_string())),
            Message::ActiveAssetCtx(active_asset_ctx) => {
                serde_json::to_string(&Subscription::ActiveAssetCtx {
                    coin: active_asset_ctx.data.coin.clone(),
                })
                .map_err(|e| Error::JsonParse(e.to_string()))
            }
            Message::ActiveSpotAssetCtx(active_spot_asset_ctx) => {
                serde_json::to_string(&Subscription::ActiveAssetCtx {
                    coin: active_spot_asset_ctx.data.coin.clone(),
                })
                .map_err(|e| Error::JsonParse(e.to_string()))
            }
            Message::ActiveAssetData(active_asset_data) => {
                serde_json::to_string(&Subscription::ActiveAssetData {
                    user: active_asset_data.data.user,
                    coin: active_asset_data.data.coin.clone(),
                })
                .map_err(|e| Error::JsonParse(e.to_string()))
            }
            Message::Bbo(bbo) => serde_json::to_string(&Subscription::Bbo {
                coin: bbo.data.coin.clone(),
            })
            .map_err(|e| Error::JsonParse(e.to_string())),
            Message::Post(_) => Ok(String::default()),
            Message::SubscriptionResponse | Message::Pong => Ok(String::default()),
            Message::NoData => Ok("".to_string()),
            Message::HyperliquidError(err) => Ok(format!("hyperliquid error: {err:?}")),
        }
    }

    async fn parse_and_send_data(
        data: std::result::Result<protocol::Message, tungstenite::Error>,
        subscriptions: &Arc<Mutex<HashMap<String, Vec<SubscriptionData>>>>,
        pending_posts: &Arc<Mutex<HashMap<u64, oneshot::Sender<ExchangeResponseStatus>>>>,
    ) -> Result<()> {
        match data {
            Ok(data) => match data.into_text() {
                Ok(data) => {
                    if !data.starts_with('{') {
                        return Ok(());
                    }
                    let message = serde_json::from_str::<Message>(&data)
                        .map_err(|e| Error::JsonParse(e.to_string()))?;

                    // Handle post responses by routing to the pending oneshot
                    if let Message::Post(ref post_envelope) = message {
                        let id = post_envelope.data.id;
                        let mut pending = pending_posts.lock().await;
                        if let Some(sender) = pending.remove(&id) {
                            let response_status = match &post_envelope.data.response {
                                PostResponsePayload::Action { payload } => payload.clone(),
                                PostResponsePayload::Info { payload: _ } => {
                                    // Not an exchange API error — drop the sender so
                                    // post_action returns Error::WsPostResponseDropped
                                    drop(sender);
                                    return Ok(());
                                }
                                PostResponsePayload::Error(err) => ExchangeResponseStatus::Err(
                                    ExchangeApiError::parse(err.payload.clone()),
                                ),
                            };
                            let _ = sender.send(response_status);
                        } else {
                            warn!("Received post response for unknown id={id}");
                        }
                        return Ok(());
                    }

                    let identifier = WsManager::get_identifier(&message)?;
                    if identifier.is_empty() {
                        return Ok(());
                    }

                    let mut subscriptions = subscriptions.lock().await;
                    let mut res = Ok(());
                    if let Some(subscription_datas) = subscriptions.get_mut(&identifier) {
                        for subscription_data in subscription_datas {
                            if let Err(e) = subscription_data
                                .sending_channel
                                .send(message.clone())
                                .map_err(|e| Error::WsSend(e.to_string()))
                            {
                                res = Err(e);
                            }
                        }
                    }
                    res
                }
                Err(err) => {
                    let error = Error::ReaderTextConversion(err.to_string());
                    Ok(WsManager::send_to_all_subscriptions(
                        subscriptions,
                        Message::HyperliquidError(error.to_string()),
                    )
                    .await?)
                }
            },
            Err(err) => {
                let error = Error::GenericReader(err.to_string());
                Ok(WsManager::send_to_all_subscriptions(
                    subscriptions,
                    Message::HyperliquidError(error.to_string()),
                )
                .await?)
            }
        }
    }

    async fn send_to_all_subscriptions(
        subscriptions: &Arc<Mutex<HashMap<String, Vec<SubscriptionData>>>>,
        message: Message,
    ) -> Result<()> {
        let mut subscriptions = subscriptions.lock().await;
        let mut res = Ok(());
        for subscription_datas in subscriptions.values_mut() {
            for subscription_data in subscription_datas {
                if let Err(e) = subscription_data
                    .sending_channel
                    .send(message.clone())
                    .map_err(|e| Error::WsSend(e.to_string()))
                {
                    res = Err(e);
                }
            }
        }
        res
    }

    fn build_subscription_payload(method: &'static str, identifier: &str) -> Result<String> {
        serde_json::to_string(&SubscriptionSendData {
            method,
            subscription: &serde_json::from_str::<serde_json::Value>(identifier)
                .map_err(|e| Error::JsonParse(e.to_string()))?,
        })
        .map_err(|e| Error::JsonParse(e.to_string()))
    }

    /// Send subscription data via the channel (normal path).
    fn send_subscription_data_via_channel(
        write_tx: &mpsc::UnboundedSender<protocol::Message>,
        method: &'static str,
        identifier: &str,
    ) -> Result<()> {
        let payload = Self::build_subscription_payload(method, identifier)?;
        write_tx
            .send(protocol::Message::Text(payload))
            .map_err(|e| Error::Websocket(e.to_string()))?;
        Ok(())
    }

    /// Send subscription data directly to the writer (used during reconnect when we hold the lock).
    async fn send_subscription_data_to_writer(
        method: &'static str,
        writer: &mut WsWriter,
        identifier: &str,
    ) -> Result<()> {
        let payload = Self::build_subscription_payload(method, identifier)?;
        writer
            .send(protocol::Message::Text(payload))
            .await
            .map_err(|e| Error::Websocket(e.to_string()))?;
        Ok(())
    }

    pub(crate) async fn add_subscription(
        &mut self,
        identifier: String,
        sending_channel: UnboundedSender<Message>,
    ) -> Result<u32> {
        let mut subscriptions = self.subscriptions.lock().await;

        let identifier_entry = if let Subscription::UserEvents { user: _ } =
            serde_json::from_str::<Subscription>(&identifier)
                .map_err(|e| Error::JsonParse(e.to_string()))?
        {
            "userEvents".to_string()
        } else if let Subscription::OrderUpdates { user: _ } =
            serde_json::from_str::<Subscription>(&identifier)
                .map_err(|e| Error::JsonParse(e.to_string()))?
        {
            "orderUpdates".to_string()
        } else {
            identifier.clone()
        };
        let subscriptions = subscriptions
            .entry(identifier_entry.clone())
            .or_insert(Vec::new());

        if !subscriptions.is_empty() && identifier_entry.eq("userEvents") {
            return Err(Error::UserEvents);
        }

        if subscriptions.is_empty() {
            Self::send_subscription_data_via_channel(&self.write_tx, "subscribe", &identifier)?;
        }

        let subscription_id = self.subscription_id;
        self.subscription_identifiers
            .insert(subscription_id, identifier.clone());
        subscriptions.push(SubscriptionData {
            sending_channel,
            subscription_id,
            id: identifier,
        });

        self.subscription_id += 1;
        Ok(subscription_id)
    }

    pub(crate) async fn remove_subscription(&mut self, subscription_id: u32) -> Result<()> {
        let identifier = self
            .subscription_identifiers
            .get(&subscription_id)
            .ok_or(Error::SubscriptionNotFound)?
            .clone();

        let identifier_entry = if let Subscription::UserEvents { user: _ } =
            serde_json::from_str::<Subscription>(&identifier)
                .map_err(|e| Error::JsonParse(e.to_string()))?
        {
            "userEvents".to_string()
        } else if let Subscription::OrderUpdates { user: _ } =
            serde_json::from_str::<Subscription>(&identifier)
                .map_err(|e| Error::JsonParse(e.to_string()))?
        {
            "orderUpdates".to_string()
        } else {
            identifier.clone()
        };

        self.subscription_identifiers.remove(&subscription_id);

        let mut subscriptions = self.subscriptions.lock().await;

        let subscriptions = subscriptions
            .get_mut(&identifier_entry)
            .ok_or(Error::SubscriptionNotFound)?;
        let index = subscriptions
            .iter()
            .position(|subscription_data| subscription_data.subscription_id == subscription_id)
            .ok_or(Error::SubscriptionNotFound)?;
        subscriptions.remove(index);

        if subscriptions.is_empty() {
            Self::send_subscription_data_via_channel(&self.write_tx, "unsubscribe", &identifier)?;
        }
        Ok(())
    }
}

impl Drop for WsManager {
    fn drop(&mut self) {
        self.stop_flag.store(true, Ordering::Relaxed);
    }
}
