use std::sync::Arc;

use axum::{
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::post,
    Json, Router,
};
use tokio::{net::TcpListener, sync::{mpsc, Mutex}};

use crate::{
    handlers::dispatch,
    types::{JsonRpcRequest, JsonRpcResponse},
};
use nu_mempool::Mempool;
use nu_state::StateDb;

pub struct AppState {
    pub db:         Arc<Mutex<StateDb>>,
    pub mempool:    Arc<Mutex<Mempool>>,
    pub chain_id:   String,
    /// Raw serialized tx bytes forwarded to P2P gossip on acceptance
    pub p2p_tx:     Option<mpsc::Sender<Vec<u8>>>,
}

pub struct RpcServer {
    pub bind_addr: String,
    pub state:     Arc<AppState>,
}

impl RpcServer {
    pub fn new(
        bind_addr: impl Into<String>,
        db:        Arc<Mutex<StateDb>>,
        mempool:   Arc<Mutex<Mempool>>,
        chain_id:  String,
    ) -> Self {
        Self::with_p2p(bind_addr, db, mempool, chain_id, None)
    }

    pub fn with_p2p(
        bind_addr: impl Into<String>,
        db:        Arc<Mutex<StateDb>>,
        mempool:   Arc<Mutex<Mempool>>,
        chain_id:  String,
        p2p_tx:    Option<mpsc::Sender<Vec<u8>>>,
    ) -> Self {
        Self {
            bind_addr: bind_addr.into(),
            state: Arc::new(AppState { db, mempool, chain_id, p2p_tx }),
        }
    }

    pub async fn run(self) -> anyhow::Result<()> {
        let router = Router::new()
            .route("/rpc", post(rpc_handler))
            .with_state(self.state);

        let listener = TcpListener::bind(&self.bind_addr).await?;
        tracing::info!("RPC listening on {}", self.bind_addr);
        axum::serve(listener, router).await?;
        Ok(())
    }
}

async fn rpc_handler(
    State(state): State<Arc<AppState>>,
    body: axum::body::Bytes,
) -> impl IntoResponse {
    let req: JsonRpcRequest = match serde_json::from_slice(&body) {
        Ok(r) => r,
        Err(_) => {
            let resp = JsonRpcResponse::err(
                serde_json::Value::Null,
                -32700,
                "Parse error".into(),
            );
            return (StatusCode::OK, Json(resp));
        }
    };

    let resp = dispatch(req, &state).await;
    (StatusCode::OK, Json(resp))
}