use axum::{
    extract::Query,
    http::{HeaderMap, StatusCode},
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    net::SocketAddr,
    sync::Arc,
    time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::sync::RwLock;
use tower_http::trace::TraceLayer;

#[derive(Clone, Debug, Serialize, Deserialize)]
struct DirectoryEntry {
    stream_id: String,
    title: String,
    offer: String,
    updated_ms: u64,
    expires_ms: u64,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct AnswerEntry {
    stream_id: String,
    answer: String,
    updated_ms: u64,
    expires_ms: u64,
}

#[derive(Clone, Debug, Serialize)]
struct DirectoryList {
    now_ms: u64,
    entries: Vec<DirectoryEntry>,
}

#[derive(Clone, Debug, Serialize)]
struct HealthResp {
    ok: bool,
}

#[derive(Clone, Debug, Deserialize)]
struct AnnounceReq {
    stream_id: String,
    title: String,
    offer: String,
    expires_ms: Option<u64>,
}

#[derive(Clone, Debug, Serialize)]
struct AnnounceResp {
    ok: bool,
    ttl_ms: u64,
    entry: DirectoryEntry,
}

#[derive(Clone, Debug, Deserialize)]
struct AnswerPostReq {
    stream_id: String,
    answer: String,
}

#[derive(Clone, Debug, Deserialize)]
struct AnswerGetReq {
    stream_id: String,
}

#[derive(Default)]
struct State {
    entries: HashMap<String, DirectoryEntry>,
    answers: HashMap<String, AnswerEntry>,
}

fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::from_secs(0))
        .as_millis() as u64
}

fn clamp_str(mut s: String, max_len: usize) -> String {
    if s.len() <= max_len {
        return s;
    }
    s.truncate(max_len);
    s
}

fn json_headers() -> HeaderMap {
    let mut headers = HeaderMap::new();
    headers.insert("content-type", "application/json; charset=utf-8".parse().unwrap());
    headers.insert("cache-control", "no-store".parse().unwrap());
    headers
}

fn prune_state(state: &mut State, now: u64) {
    state.entries.retain(|_, v| v.expires_ms > now);
    state.answers.retain(|_, v| v.expires_ms > now);

    // Cap growth defensively. This is not spam-resistant; it's a bootstrap rendezvous.
    if state.entries.len() > 200 {
        let mut items = state.entries.values().cloned().collect::<Vec<_>>();
        items.sort_by_key(|e| std::cmp::Reverse(e.updated_ms));
        items.truncate(200);
        state.entries = items.into_iter().map(|e| (e.stream_id.clone(), e)).collect();
    }
    if state.answers.len() > 500 {
        let mut items = state.answers.values().cloned().collect::<Vec<_>>();
        items.sort_by_key(|e| std::cmp::Reverse(e.updated_ms));
        items.truncate(500);
        state.answers = items.into_iter().map(|e| (e.stream_id.clone(), e)).collect();
    }
}

async fn health() -> impl IntoResponse {
    (json_headers(), Json(HealthResp { ok: true }))
}

async fn directory(state: axum::extract::State<Arc<RwLock<State>>>) -> impl IntoResponse {
    let now = now_ms();
    let mut guard = state.write().await;
    prune_state(&mut guard, now);
    let mut entries = guard.entries.values().cloned().collect::<Vec<_>>();
    entries.sort_by_key(|e| std::cmp::Reverse(e.updated_ms));
    (json_headers(), Json(DirectoryList { now_ms: now, entries }))
}

async fn announce(
    state: axum::extract::State<Arc<RwLock<State>>>,
    Json(body): Json<AnnounceReq>,
) -> impl IntoResponse {
    let now = now_ms();

    if body.stream_id.is_empty() || body.title.is_empty() || body.offer.is_empty() {
        let resp = serde_json::json!({ "error": "missing stream_id/title/offer" });
        return (StatusCode::BAD_REQUEST, json_headers(), Json(resp)).into_response();
    }
    if body.offer.len() > 64_000 {
        let resp = serde_json::json!({ "error": "offer too large" });
        return (StatusCode::PAYLOAD_TOO_LARGE, json_headers(), Json(resp)).into_response();
    }

    let requested_expires = body.expires_ms.unwrap_or(now + 20_000);
    let requested_ttl = requested_expires.saturating_sub(now);
    let ttl_ms = requested_ttl.clamp(5_000, 60_000);

    let entry = DirectoryEntry {
        stream_id: clamp_str(body.stream_id, 256),
        title: clamp_str(body.title, 128),
        offer: body.offer,
        updated_ms: now,
        expires_ms: now + ttl_ms,
    };

    let mut guard = state.write().await;
    prune_state(&mut guard, now);
    guard.entries.insert(entry.stream_id.clone(), entry.clone());

    (
        json_headers(),
        Json(AnnounceResp {
            ok: true,
            ttl_ms,
            entry,
        }),
    )
        .into_response()
}

async fn post_answer(
    state: axum::extract::State<Arc<RwLock<State>>>,
    Json(body): Json<AnswerPostReq>,
) -> impl IntoResponse {
    let now = now_ms();
    if body.stream_id.is_empty() || body.answer.is_empty() {
        let resp = serde_json::json!({ "error": "missing stream_id/answer" });
        return (StatusCode::BAD_REQUEST, json_headers(), Json(resp)).into_response();
    }
    if body.answer.len() > 64_000 {
        let resp = serde_json::json!({ "error": "answer too large" });
        return (StatusCode::PAYLOAD_TOO_LARGE, json_headers(), Json(resp)).into_response();
    }

    let entry = AnswerEntry {
        stream_id: clamp_str(body.stream_id, 256),
        answer: body.answer,
        updated_ms: now,
        expires_ms: now + 2 * 60_000,
    };

    let mut guard = state.write().await;
    prune_state(&mut guard, now);
    guard.answers.insert(entry.stream_id.clone(), entry);

    (json_headers(), Json(serde_json::json!({ "ok": true }))).into_response()
}

async fn get_answer(
    state: axum::extract::State<Arc<RwLock<State>>>,
    Query(q): Query<AnswerGetReq>,
) -> impl IntoResponse {
    let now = now_ms();
    if q.stream_id.is_empty() {
        let resp = serde_json::json!({ "error": "missing stream_id" });
        return (StatusCode::BAD_REQUEST, json_headers(), Json(resp)).into_response();
    }

    let mut guard = state.write().await;
    prune_state(&mut guard, now);

    // One-shot: first reader consumes.
    let Some(answer) = guard.answers.remove(&q.stream_id) else {
        let resp = serde_json::json!({ "error": "not found" });
        return (StatusCode::NOT_FOUND, json_headers(), Json(resp)).into_response();
    };

    (json_headers(), Json(answer)).into_response()
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            std::env::var("RUST_LOG").unwrap_or_else(|_| "info,tower_http=info".to_string()),
        )
        .init();

    let state = Arc::new(RwLock::new(State::default()));
    let app = Router::new()
        .route("/api/health", get(health))
        .route("/api/directory", get(directory))
        .route("/api/announce", post(announce))
        .route("/api/answer", post(post_answer).get(get_answer))
        .with_state(state)
        .layer(TraceLayer::new_for_http());

    let addr: SocketAddr = "0.0.0.0:8080".parse().unwrap();
    tracing::info!("ec-cf-bootstrap-api listening on {}", addr);
    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;
    Ok(())
}

async fn shutdown_signal() {
    let _ = tokio::signal::ctrl_c().await;
}
