// SquirrelDB Rust SDK - Client
// Generated by sdk-generator v0.1.0
// DO NOT EDIT MANUALLY

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

use futures::Stream;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::TcpStream;
use tokio::sync::{mpsc, oneshot, RwLock};
use uuid::Uuid;

use crate::error::{Error, Result};
use crate::protocol::*;

/// Options for connecting to SquirrelDB
#[derive(Debug, Clone, Default)]
pub struct ConnectOptions {
    pub host: String,
    pub port: u16,
    pub auth_token: Option<String>,
    pub use_messagepack: bool,
    pub json_fallback: bool,
}

impl ConnectOptions {
    pub fn new(host: impl Into<String>, port: u16) -> Self {
        Self {
            host: host.into(),
            port,
            auth_token: None,
            use_messagepack: true,
            json_fallback: true,
        }
    }

    pub fn with_auth(mut self, token: impl Into<String>) -> Self {
        self.auth_token = Some(token.into());
        self
    }
}

#[derive(Debug)]
struct PendingRequest {
    tx: oneshot::Sender<Result<ServerMessage>>,
}

#[derive(Debug)]
struct ActiveSubscription {
    tx: mpsc::UnboundedSender<ChangeEvent>,
}

/// SquirrelDB client
#[derive(Debug)]
pub struct SquirrelDB {
    session_id: Uuid,
    encoding: Encoding,
    writer: Arc<tokio::sync::Mutex<BufWriter<OwnedWriteHalf>>>,
    pending: Arc<RwLock<HashMap<String, PendingRequest>>>,
    subscriptions: Arc<RwLock<HashMap<String, ActiveSubscription>>>,
    request_id: AtomicU64,
    _reader_task: tokio::task::JoinHandle<()>,
}

impl SquirrelDB {
    /// Connect to SquirrelDB server
    pub async fn connect(addr: &str) -> Result<Self> {
        let parts: Vec<&str> = addr.split(':').collect();
        let host = parts.first().unwrap_or(&"localhost");
        let port: u16 = parts.get(1).and_then(|p| p.parse().ok()).unwrap_or(8082);
        Self::connect_with_options(ConnectOptions::new(*host, port)).await
    }

    /// Connect with full options
    pub async fn connect_with_options(opts: ConnectOptions) -> Result<Self> {
        let addr = format!("{}:{}", opts.host, opts.port);
        let mut stream = TcpStream::connect(&addr)
            .await
            .map_err(|e| Error::Connection(e.to_string()))?;

        let (session_id, encoding) = Self::handshake(&mut stream, &opts).await?;
        let (read_half, write_half) = stream.into_split();
        let reader = BufReader::new(read_half);
        let writer = Arc::new(tokio::sync::Mutex::new(BufWriter::new(write_half)));

        let pending: Arc<RwLock<HashMap<String, PendingRequest>>> =
            Arc::new(RwLock::new(HashMap::new()));
        let subscriptions: Arc<RwLock<HashMap<String, ActiveSubscription>>> =
            Arc::new(RwLock::new(HashMap::new()));

        let reader_task = tokio::spawn(Self::reader_loop(
            reader,
            encoding,
            pending.clone(),
            subscriptions.clone(),
        ));

        Ok(Self {
            session_id,
            encoding,
            writer,
            pending,
            subscriptions,
            request_id: AtomicU64::new(1),
            _reader_task: reader_task,
        })
    }

    pub fn session_id(&self) -> Uuid {
        self.session_id
    }

    async fn handshake(stream: &mut TcpStream, opts: &ConnectOptions) -> Result<(Uuid, Encoding)> {
        stream.write_all(MAGIC).await?;
        stream.write_u8(PROTOCOL_VERSION).await?;

        let flags = ProtocolFlags {
            messagepack: opts.use_messagepack,
            json_fallback: opts.json_fallback,
        };
        stream.write_u8(flags.into()).await?;

        let token = opts.auth_token.as_deref().unwrap_or("");
        let token_bytes = token.as_bytes();
        stream.write_u16(token_bytes.len() as u16).await?;
        if !token_bytes.is_empty() {
            stream.write_all(token_bytes).await?;
        }
        stream.flush().await?;

        let status_byte = stream.read_u8().await?;
        let status = HandshakeStatus::try_from(status_byte)
            .map_err(|_| Error::Handshake(format!("Invalid status: {}", status_byte)))?;

        let server_version = stream.read_u8().await?;
        let server_flags_byte = stream.read_u8().await?;

        let mut session_bytes = [0u8; 16];
        stream.read_exact(&mut session_bytes).await?;
        let session_id = Uuid::from_bytes(session_bytes);

        match status {
            HandshakeStatus::Success => {}
            HandshakeStatus::VersionMismatch => {
                return Err(Error::VersionMismatch {
                    server: server_version,
                    client: PROTOCOL_VERSION,
                });
            }
            HandshakeStatus::AuthFailed => {
                return Err(Error::AuthFailed);
            }
        }

        let server_flags = ProtocolFlags::from(server_flags_byte);
        let encoding = if server_flags.messagepack {
            Encoding::MessagePack
        } else {
            Encoding::Json
        };

        Ok((session_id, encoding))
    }

    async fn reader_loop(
        mut reader: BufReader<OwnedReadHalf>,
        encoding: Encoding,
        pending: Arc<RwLock<HashMap<String, PendingRequest>>>,
        subscriptions: Arc<RwLock<HashMap<String, ActiveSubscription>>>,
    ) {
        loop {
            match Self::read_frame(&mut reader, encoding).await {
                Ok(msg) => {
                    Self::dispatch_message(msg, &pending, &subscriptions).await;
                }
                Err(_) => break,
            }
        }
    }

    async fn dispatch_message(
        msg: ServerMessage,
        pending: &Arc<RwLock<HashMap<String, PendingRequest>>>,
        subscriptions: &Arc<RwLock<HashMap<String, ActiveSubscription>>>,
    ) {
        match &msg {
            ServerMessage::Change { id, change } => {
                if let Some(sub) = subscriptions.read().await.get(id) {
                    let _ = sub.tx.send(change.clone());
                }
            }
            ServerMessage::Result { id, .. }
            | ServerMessage::Subscribed { id }
            | ServerMessage::Unsubscribed { id }
            | ServerMessage::Error { id, .. }
            | ServerMessage::Pong { id } => {
                if let Some(req) = pending.write().await.remove(id) {
                    let _ = req.tx.send(Ok(msg));
                }
            }
        }
    }

    async fn read_frame(
        reader: &mut BufReader<OwnedReadHalf>,
        encoding: Encoding,
    ) -> Result<ServerMessage> {
        let length = reader.read_u32().await?;
        if length > MAX_MESSAGE_SIZE {
            return Err(Error::Server(format!("Message too large: {}", length)));
        }

        let _msg_type = reader.read_u8().await?;
        let _encoding = reader.read_u8().await?;

        let payload_len = length as usize - 2;
        let mut payload = vec![0u8; payload_len];
        reader.read_exact(&mut payload).await?;

        let msg: ServerMessage = match encoding {
            Encoding::MessagePack => rmp_serde::from_slice(&payload)?,
            Encoding::Json => serde_json::from_slice(&payload)?,
        };

        Ok(msg)
    }

    async fn write_frame(&self, msg: &ClientMessage) -> Result<()> {
        let payload = match self.encoding {
            Encoding::MessagePack => rmp_serde::to_vec(msg)?,
            Encoding::Json => serde_json::to_vec(msg)?,
        };

        let length = (payload.len() + 2) as u32;
        let mut writer = self.writer.lock().await;
        writer.write_u32(length).await?;
        writer.write_u8(MessageType::Request as u8).await?;
        writer.write_u8(self.encoding as u8).await?;
        writer.write_all(&payload).await?;
        writer.flush().await?;

        Ok(())
    }

    fn next_id(&self) -> String {
        self.request_id.fetch_add(1, Ordering::SeqCst).to_string()
    }

    async fn request(&self, msg: ClientMessage) -> Result<ServerMessage> {
        let id = match &msg {
            ClientMessage::Query { id, .. }
            | ClientMessage::Subscribe { id, .. }
            | ClientMessage::Unsubscribe { id }
            | ClientMessage::Insert { id, .. }
            | ClientMessage::Update { id, .. }
            | ClientMessage::Delete { id, .. }
            | ClientMessage::ListCollections { id }
            | ClientMessage::Ping { id } => id.clone(),
        };

        let (tx, rx) = oneshot::channel();
        self.pending.write().await.insert(id.clone(), PendingRequest { tx });

        self.write_frame(&msg).await?;

        rx.await.map_err(|_| Error::ChannelClosed)?
    }

    /// Execute a query
    pub async fn query<T: serde::de::DeserializeOwned>(&self, query: &str) -> Result<T> {
        let msg = ClientMessage::Query {
            id: self.next_id(),
            query: query.to_string(),
        };

        match self.request(msg).await? {
            ServerMessage::Result { data, .. } => {
                serde_json::from_value(data).map_err(|e| Error::Serialization(e.to_string()))
            }
            ServerMessage::Error { error, .. } => Err(Error::Server(error)),
            _ => Err(Error::Server("Unexpected response".to_string())),
        }
    }

    /// Insert a document
    pub async fn insert(&self, collection: &str, data: serde_json::Value) -> Result<Document> {
        let msg = ClientMessage::Insert {
            id: self.next_id(),
            collection: collection.to_string(),
            data,
        };

        match self.request(msg).await? {
            ServerMessage::Result { data, .. } => {
                serde_json::from_value(data).map_err(|e| Error::Serialization(e.to_string()))
            }
            ServerMessage::Error { error, .. } => Err(Error::Server(error)),
            _ => Err(Error::Server("Unexpected response".to_string())),
        }
    }

    /// Update a document
    pub async fn update(&self, collection: &str, document_id: Uuid, data: serde_json::Value) -> Result<Document> {
        let msg = ClientMessage::Update {
            id: self.next_id(),
            collection: collection.to_string(),
            document_id,
            data,
        };

        match self.request(msg).await? {
            ServerMessage::Result { data, .. } => {
                serde_json::from_value(data).map_err(|e| Error::Serialization(e.to_string()))
            }
            ServerMessage::Error { error, .. } => Err(Error::Server(error)),
            _ => Err(Error::Server("Unexpected response".to_string())),
        }
    }

    /// Delete a document
    pub async fn delete(&self, collection: &str, document_id: Uuid) -> Result<Document> {
        let msg = ClientMessage::Delete {
            id: self.next_id(),
            collection: collection.to_string(),
            document_id,
        };

        match self.request(msg).await? {
            ServerMessage::Result { data, .. } => {
                serde_json::from_value(data).map_err(|e| Error::Serialization(e.to_string()))
            }
            ServerMessage::Error { error, .. } => Err(Error::Server(error)),
            _ => Err(Error::Server("Unexpected response".to_string())),
        }
    }

    /// List all collections
    pub async fn list_collections(&self) -> Result<Vec<String>> {
        let msg = ClientMessage::ListCollections { id: self.next_id() };

        match self.request(msg).await? {
            ServerMessage::Result { data, .. } => {
                serde_json::from_value(data).map_err(|e| Error::Serialization(e.to_string()))
            }
            ServerMessage::Error { error, .. } => Err(Error::Server(error)),
            _ => Err(Error::Server("Unexpected response".to_string())),
        }
    }

    /// Subscribe to changes
    pub async fn subscribe(&self, query: &str) -> Result<Subscription<'_>> {
        let id = self.next_id();
        let msg = ClientMessage::Subscribe {
            id: id.clone(),
            query: query.to_string(),
        };

        let (tx, rx) = mpsc::unbounded_channel();
        self.subscriptions.write().await.insert(id.clone(), ActiveSubscription { tx });

        match self.request(msg).await? {
            ServerMessage::Subscribed { .. } => Ok(Subscription {
                id,
                rx,
                client_subs: self.subscriptions.clone(),
                client_writer: self.writer.clone(),
                encoding: self.encoding,
                request_id: &self.request_id,
            }),
            ServerMessage::Error { error, .. } => {
                self.subscriptions.write().await.remove(&id);
                Err(Error::Server(error))
            }
            _ => {
                self.subscriptions.write().await.remove(&id);
                Err(Error::Server("Unexpected response".to_string()))
            }
        }
    }

    /// Ping the server
    pub async fn ping(&self) -> Result<()> {
        let msg = ClientMessage::Ping { id: self.next_id() };

        match self.request(msg).await? {
            ServerMessage::Pong { .. } => Ok(()),
            ServerMessage::Error { error, .. } => Err(Error::Server(error)),
            _ => Err(Error::Server("Unexpected response".to_string())),
        }
    }
}

/// Active subscription handle
pub struct Subscription<'a> {
    id: String,
    rx: mpsc::UnboundedReceiver<ChangeEvent>,
    client_subs: Arc<RwLock<HashMap<String, ActiveSubscription>>>,
    client_writer: Arc<tokio::sync::Mutex<BufWriter<OwnedWriteHalf>>>,
    encoding: Encoding,
    request_id: &'a AtomicU64,
}

impl<'a> Subscription<'a> {
    pub async fn next(&mut self) -> Option<ChangeEvent> {
        self.rx.recv().await
    }

    pub async fn unsubscribe(self) -> Result<()> {
        self.client_subs.write().await.remove(&self.id);
        let id = self.request_id.fetch_add(1, Ordering::SeqCst).to_string();
        let msg = ClientMessage::Unsubscribe { id };

        let payload = match self.encoding {
            Encoding::MessagePack => rmp_serde::to_vec(&msg)?,
            Encoding::Json => serde_json::to_vec(&msg)?,
        };

        let length = (payload.len() + 2) as u32;
        let mut writer = self.client_writer.lock().await;
        writer.write_u32(length).await?;
        writer.write_u8(MessageType::Request as u8).await?;
        writer.write_u8(self.encoding as u8).await?;
        writer.write_all(&payload).await?;
        writer.flush().await?;

        Ok(())
    }
}

impl Stream for Subscription<'_> {
    type Item = ChangeEvent;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        std::pin::Pin::new(&mut self.rx).poll_recv(cx)
    }
}
