mod proxy;
mod rpc;
mod ws;

use std::sync::Arc;

use actix_cors::Cors;
use actix_web::{web, App, HttpRequest, HttpResponse, HttpServer};
use clap::Parser;
use rust_embed::Embed;
use serde_json::Value;

use rpc::{classify_method, JsonRpcRequest, JsonRpcResponse, MethodType};
use ws::AppState;

#[derive(Embed)]
#[folder = "frontend/out"]
struct FrontendAssets;

#[derive(Parser)]
#[command(name = "prpc", about = "Proxy RPC — routes signing requests to a frontend wallet")]
struct Cli {
    /// Upstream RPC URL to forward read calls to
    #[arg(long)]
    rpc_url: String,

    /// Port for JSON-RPC proxy + frontend UI
    #[arg(long, default_value = "8545")]
    port: u16,

    /// Port for the frontend UI (if you want it on a separate port)
    #[arg(long)]
    ui_port: Option<u16>,

    /// Reown (WalletConnect) project ID from https://cloud.reown.com
    #[arg(long, env = "REOWN_PROJECT_ID")]
    reown_project_id: String,
}

/// Runtime config shared across handlers.
pub struct Config {
    pub reown_project_id: String,
    pub ws_url: String,
}

async fn handle_rpc(
    state: web::Data<Arc<AppState>>,
    body: web::Json<Value>,
) -> HttpResponse {
    let raw = body.into_inner();
    log::debug!("Incoming RPC: {}", serde_json::to_string(&raw).unwrap_or_default());

    // Handle batch requests
    if let Value::Array(batch) = &raw {
        let requests: Vec<JsonRpcRequest> = match batch
            .iter()
            .map(|v| serde_json::from_value(v.clone()))
            .collect::<Result<Vec<_>, _>>()
        {
            Ok(r) => r,
            Err(e) => {
                return HttpResponse::Ok().json(JsonRpcResponse::error(
                    Value::Null,
                    -32700,
                    format!("Parse error: {e}"),
                ));
            }
        };

        let mut responses = Vec::with_capacity(requests.len());
        for req in &requests {
            responses.push(handle_single_request(&state, req).await);
        }
        return HttpResponse::Ok().json(responses);
    }

    // Single request
    let request: JsonRpcRequest = match serde_json::from_value(raw) {
        Ok(r) => r,
        Err(e) => {
            return HttpResponse::Ok().json(JsonRpcResponse::error(
                Value::Null,
                -32700,
                format!("Parse error: {e}"),
            ));
        }
    };

    let response = handle_single_request(&state, &request).await;
    HttpResponse::Ok().json(response)
}

async fn handle_single_request(
    state: &AppState,
    request: &JsonRpcRequest,
) -> JsonRpcResponse {
    let method_type = classify_method(&request.method);

    match method_type {
        MethodType::Read => {
            log::debug!("Forwarding read call: {}", request.method);
            proxy::forward_to_upstream(&state.http_client, &state.rpc_url, request).await
        }
        MethodType::Account => {
            let addr = state.connected_address.read().await;
            match addr.as_ref() {
                Some(address) => {
                    log::debug!("Returning connected account: {address}");
                    JsonRpcResponse::success(
                        request.id.clone(),
                        Value::Array(vec![Value::String(address.clone())]),
                    )
                }
                None => {
                    log::debug!("No wallet connected, returning empty accounts");
                    JsonRpcResponse::success(request.id.clone(), Value::Array(vec![]))
                }
            }
        }
        MethodType::Write => {
            log::info!("Intercepting write call: {}", request.method);

            let request_id = uuid::Uuid::new_v4().to_string();

            let rx = state
                .send_signing_request(
                    request_id.clone(),
                    request.method.clone(),
                    request.params.clone(),
                )
                .await;

            let rx = match rx {
                Some(rx) => rx,
                None => {
                    return JsonRpcResponse::error(
                        request.id.clone(),
                        -32603,
                        "No frontend connected. Open the signing UI and connect a wallet.",
                    );
                }
            };

            // Wait for the frontend to respond (120s timeout)
            match tokio::time::timeout(std::time::Duration::from_secs(120), rx).await {
                Ok(Ok(mut resp)) => {
                    resp.id = request.id.clone();
                    resp
                }
                Ok(Err(_)) => JsonRpcResponse::error(
                    request.id.clone(),
                    -32603,
                    "Frontend connection lost while waiting for signature",
                ),
                Err(_) => {
                    state.pending.remove(&request_id);
                    JsonRpcResponse::error(
                        request.id.clone(),
                        -32603,
                        "Signing request timed out (120s)",
                    )
                }
            }
        }
    }
}

async fn ws_handler(
    req: HttpRequest,
    stream: web::Payload,
    state: web::Data<Arc<AppState>>,
) -> Result<HttpResponse, actix_web::Error> {
    let (response, session, msg_stream) = actix_ws::handle(&req, stream)?;

    let state = state.get_ref().clone();
    actix_web::rt::spawn(ws::handle_ws_connection(state, session, msg_stream));

    Ok(response)
}

/// Serves /api/config — runtime config for the frontend.
async fn api_config(config: web::Data<Arc<Config>>) -> HttpResponse {
    HttpResponse::Ok().json(serde_json::json!({
        "projectId": config.reown_project_id,
        "wsUrl": config.ws_url,
    }))
}

/// Serves embedded frontend static files.
async fn frontend_handler(req: HttpRequest) -> HttpResponse {
    let path = req.path().trim_start_matches('/');
    let path = if path.is_empty() { "index.html" } else { path };

    // Try exact path first
    if let Some(file) = FrontendAssets::get(path) {
        let mime = mime_guess::from_path(path).first_or_octet_stream();
        return HttpResponse::Ok()
            .content_type(mime.as_ref())
            .body(file.data.into_owned());
    }

    // For paths without extension (SPA routes), serve index.html
    if !path.contains('.') {
        if let Some(file) = FrontendAssets::get("index.html") {
            return HttpResponse::Ok()
                .content_type("text/html")
                .body(file.data.into_owned());
        }
    }

    HttpResponse::NotFound().body("Not found")
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let cli = Cli::parse();
    let ui_port = cli.ui_port.unwrap_or(cli.port);
    let separate_ui = cli.ui_port.is_some();

    log::info!("Starting prpc");
    log::info!("  Upstream RPC: {}", cli.rpc_url);
    log::info!("  Proxy:        http://0.0.0.0:{}", cli.port);
    if separate_ui {
        log::info!("  Frontend UI:  http://0.0.0.0:{ui_port}");
    } else {
        log::info!("  Frontend UI:  http://0.0.0.0:{}", cli.port);
    }

    let state = Arc::new(AppState::new(cli.rpc_url));

    // Fetch chain ID from upstream
    if let Some(chain_id_val) = proxy::fetch_chain_id(&state.http_client, &state.rpc_url).await {
        if let Some(cid) = chain_id_val.as_str() {
            log::info!("  Chain ID:     {cid}");
            let mut lock = state.chain_id.write().await;
            *lock = Some(cid.to_string());
        }
    } else {
        log::warn!("Could not fetch chain ID from upstream RPC");
    }

    // Always return explicit WS URL so it works regardless of how the frontend is accessed
    let ws_url = format!("ws://localhost:{}/ws", cli.port);

    let config = Arc::new(Config {
        reown_project_id: cli.reown_project_id,
        ws_url,
    });

    if separate_ui {
        // Two separate servers: proxy on --port, UI on --ui-port
        let state2 = state.clone();
        let config2 = config.clone();

        let proxy_server = HttpServer::new(move || {
            App::new()
                .wrap(Cors::permissive())
                .app_data(web::Data::new(state.clone()))
                .app_data(web::Data::new(config.clone()))
                .route("/", web::post().to(handle_rpc))
                .route("/ws", web::get().to(ws_handler))
                .route("/api/config", web::get().to(api_config))
        })
        .bind(("0.0.0.0", cli.port))?
        .run();

        let ui_server = HttpServer::new(move || {
            App::new()
                .wrap(Cors::permissive())
                .app_data(web::Data::new(state2.clone()))
                .app_data(web::Data::new(config2.clone()))
                .route("/api/config", web::get().to(api_config))
                .default_service(web::to(frontend_handler))
        })
        .bind(("0.0.0.0", ui_port))?
        .run();

        tokio::try_join!(proxy_server, ui_server)?;
    } else {
        // Single server: everything on one port
        // POST / = JSON-RPC, GET / = frontend UI, GET /ws = WebSocket
        HttpServer::new(move || {
            App::new()
                .wrap(Cors::permissive())
                .app_data(web::Data::new(state.clone()))
                .app_data(web::Data::new(config.clone()))
                .route("/", web::post().to(handle_rpc))
                .route("/ws", web::get().to(ws_handler))
                .route("/api/config", web::get().to(api_config))
                // All other GET requests serve frontend static files
                .default_service(web::get().to(frontend_handler))
        })
        .bind(("0.0.0.0", cli.port))?
        .run()
        .await?;
    }

    Ok(())
}
