use crate::base::*;
use axon_core::*;
use futures_util::{SinkExt, StreamExt};
use serde::Serializer;
use std::fmt;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio_tungstenite::accept_hdr_async;
use tokio_tungstenite::tungstenite::Message;

#[derive(Debug)]
pub enum ServerError {
    Io(std::io::Error),
    AddrParse(std::net::AddrParseError),
    WebSocketError(tokio_tungstenite::tungstenite::Error),
}

impl fmt::Display for ServerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ServerError::Io(e) => write!(f, "IO error: {}", e),
            ServerError::AddrParse(e) => write!(f, "Address parse error: {}", e),
            ServerError::WebSocketError(e) => write!(f, "WebSocket error: {}", e),
        }
    }
}

impl std::error::Error for ServerError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            ServerError::Io(e) => Some(e),
            ServerError::AddrParse(e) => Some(e),
            ServerError::WebSocketError(e) => Some(e),
        }
    }
}

impl From<std::io::Error> for ServerError {
    fn from(e: std::io::Error) -> Self {
        ServerError::Io(e)
    }
}

impl From<std::net::AddrParseError> for ServerError {
    fn from(e: std::net::AddrParseError) -> Self {
        ServerError::AddrParse(e)
    }
}

impl From<tokio_tungstenite::tungstenite::Error> for ServerError {
    fn from(e: tokio_tungstenite::tungstenite::Error) -> Self {
        ServerError::WebSocketError(e)
    }
}

type Tx = mpsc::UnboundedSender<Vec<u8>>;

struct ClientConnection {
    #[allow(dead_code)]
    client_id: u64,
    tx: Tx,
}

pub struct Server {
    clients: Arc<RwLock<std::collections::HashMap<u64, ClientConnection>>>,
    snapshot: crate::base::AxonSnapshot,
    count: u32,
    pending_events: Arc<Mutex<Vec<AxonReceiveEvent>>>,
}

impl Server {
    pub fn new(addr: &str) -> Result<Self, ServerError> {
        let clients: Arc<RwLock<std::collections::HashMap<u64, ClientConnection>>> =
            Arc::new(RwLock::new(std::collections::HashMap::new()));
        let pending_events = Arc::new(Mutex::new(Vec::new()));

        let server = Self {
            clients: clients.clone(),
            snapshot: crate::base::AxonSnapshot::default(),
            count: 0,
            pending_events: pending_events.clone(),
        };

        let addr = addr.to_string();
        std::thread::spawn(move || {
            let rt = tokio::runtime::Runtime::new().expect("Failed to create Tokio runtime");
            rt.block_on(async move {
                let listener = tokio::net::TcpListener::bind(&addr)
                    .await
                    .expect("Failed to bind");
                loop {
                    match listener.accept().await {
                        Ok((stream, addr)) => {
                            let clients = clients.clone();
                            let pending = pending_events.clone();
                            tokio::spawn(async move {
                                if let Err(e) =
                                    handle_connection(stream, addr, clients, pending).await
                                {
                                    eprintln!("Connection error: {}", e);
                                }
                            });
                        }
                        Err(e) => {
                            eprintln!("Accept error: {}", e);
                        }
                    }
                }
            });
        });

        Ok(server)
    }
}

async fn handle_connection(
    stream: TcpStream,
    _addr: SocketAddr,
    clients: Arc<RwLock<std::collections::HashMap<u64, ClientConnection>>>,
    pending_events: Arc<Mutex<Vec<AxonReceiveEvent>>>,
) -> Result<(), ServerError> {
    let mut client_id = rand::random::<u64>();
    let ws_stream = accept_hdr_async(
        stream,
        |req: &tokio_tungstenite::tungstenite::handshake::server::Request,
         response: tokio_tungstenite::tungstenite::handshake::server::Response| {
            let uri = req.uri().to_string();
            client_id = parse_client_id_from_url(&uri);
            Ok(response)
        },
    )
    .await?;
    let (mut ws_sender, mut ws_receiver) = ws_stream.split();

    let (tx, mut rx) = mpsc::unbounded_channel::<Vec<u8>>();
    clients.write().await.insert(
        client_id,
        ClientConnection {
            client_id,
            tx: tx.clone(),
        },
    );

    pending_events
        .lock()
        .await
        .push(AxonReceiveEvent::AddClient { client_id });

    let client_id_for_recv = client_id;
    let pending_for_recv = pending_events.clone();

    let recv_task = tokio::spawn(async move {
        while let Some(msg) = ws_receiver.next().await {
            match msg {
                Ok(Message::Binary(data)) => {
                    let mut events = Vec::new();
                    crate::base::receive_events(client_id_for_recv, &data, &mut events);
                    let mut pending = pending_events.lock().await;
                    pending.extend(events);
                }
                Ok(Message::Close(_)) => break,
                Err(_) => break,
                _ => {}
            }
        }
        pending_for_recv
            .lock()
            .await
            .push(AxonReceiveEvent::RemoveClient {
                client_id: client_id_for_recv,
            });
    });

    let send_task = tokio::spawn(async move {
        while let Some(data) = rx.recv().await {
            if ws_sender.send(Message::Binary(data.into())).await.is_err() {
                break;
            }
        }
    });

    tokio::select! {
        _ = recv_task => {},
        _ = send_task => {},
    }

    clients.write().await.remove(&client_id);

    Ok(())
}

fn parse_client_id_from_url(url: &str) -> u64 {
    if let Some(query_start) = url.find('?') {
        let query = &url[query_start + 1..];
        for param in query.split('&') {
            if let Some(eq_pos) = param.find('=') {
                let key = &param[..eq_pos];
                let value = &param[eq_pos + 1..];
                if key == "id" {
                    if let Ok(id) = value.parse::<u64>() {
                        return id;
                    }
                }
            }
        }
    }
    rand::random::<u64>()
}

impl IAxonContext for Server {
    fn spawn(&mut self, id: u64, t: u32) {
        self.snapshot.spawn(id, t);
        let mut s = axon_sbin::SbinSerializer::new();
        s.serialize_u8(ACTION_TYPE_SPAWN).unwrap();
        s.serialize_u64(id).unwrap();
        s.serialize_u32(t).unwrap();
        s.serialize_bytes(&[]).unwrap();

        let data = s.into_vec();
        let clients = self.clients.try_write().unwrap();
        for client in clients.values() {
            let _ = client.tx.send(data.clone());
        }
        self.count += 1;
        println!("spawn: {} {}", id, t);
    }

    fn despawn(&mut self, id: u64) {
        self.snapshot.despawn(id);
        let mut s = axon_sbin::SbinSerializer::new();
        s.serialize_u8(ACTION_TYPE_DESPAWN).unwrap();
        s.serialize_u64(id).unwrap();
        s.serialize_u32(0).unwrap();
        s.serialize_bytes(&[]).unwrap();

        let data = s.into_vec();
        let clients = self.clients.try_write().unwrap();
        for client in clients.values() {
            let _ = client.tx.send(data.clone());
        }
        self.count += 1;
        println!("despawn: {}", id);
    }

    fn change(&mut self, id: u64, t: u32, v: &[u8]) {
        self.snapshot.change(id, t, v);
        let mut s = axon_sbin::SbinSerializer::new();
        s.serialize_u8(ACTION_TYPE_CHANGE).unwrap();
        s.serialize_u64(id).unwrap();
        s.serialize_u32(t).unwrap();
        s.serialize_bytes(v).unwrap();

        let data = s.into_vec();
        let clients = self.clients.try_write().unwrap();
        for client in clients.values() {
            let _ = client.tx.send(data.clone());
        }
        self.count += 1;
    }

    fn invoke(&mut self, client_id: u64, id: u64, t: u32, v: &[u8]) {
        let mut s = axon_sbin::SbinSerializer::new();
        s.serialize_u8(ACTION_TYPE_INVOKE).unwrap();
        s.serialize_u64(id).unwrap();
        s.serialize_u32(t).unwrap();
        s.serialize_bytes(v).unwrap();

        let data = s.into_vec();
        let clients = self.clients.try_write().unwrap();
        if let Some(client) = clients.get(&client_id) {
            let _ = client.tx.send(data);
        }
        self.count += 1;
    }

    fn broadcast(&mut self, id: u64, t: u32, v: &[u8]) {
        let mut s = axon_sbin::SbinSerializer::new();
        s.serialize_u8(ACTION_TYPE_INVOKE).unwrap();
        s.serialize_u64(id).unwrap();
        s.serialize_u32(t).unwrap();
        s.serialize_bytes(v).unwrap();

        let data = s.into_vec();
        let clients = self.clients.try_write().unwrap();
        for client in clients.values() {
            let _ = client.tx.send(data.clone());
        }
        self.count += 1;
    }

    fn update(&mut self, _: std::time::Duration) {}

    fn receive_event(&mut self, events: &mut Vec<AxonReceiveEvent>) {
        let mut i = events.len();
        if let Ok(mut pending) = self.pending_events.try_lock() {
            events.extend(pending.drain(..));
        }
        while i < events.len() {
            let event = &events[i];
            if let AxonReceiveEvent::AddClient { client_id } = event {
                let s = self.snapshot.snapshot();
                let clients = self.clients.try_write().unwrap();
                if let Some(client) = clients.get(&client_id) {
                    let _ = client.tx.send(s);
                }
                self.count += 1;
            }
            i += 1;
        }
    }

    fn flush(&mut self) {
        if self.count > 0 {
            self.count = 0;
            let mut s = axon_sbin::SbinSerializer::new();
            s.serialize_u8(ACTION_TYPE_FRAME).unwrap();
            s.serialize_u64(0).unwrap();
            s.serialize_u32(0).unwrap();
            s.serialize_bytes(&[]).unwrap();
            let data = s.into_vec();
            let clients = self.clients.try_write().unwrap();
            for client in clients.values() {
                let _ = client.tx.send(data.clone());
            }
        }
    }
}
