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, Deserialize)]
struct PublicStreamRelay {
    relay_url: String,
    broadcast_name: String,
    track_name: String,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct PublicStreamEntry {
    stream_id: String,
    title: String,
    relay_url: String,
    broadcast_name: String,
    track_name: String,
    relays: Vec<PublicStreamRelay>,
    updated_ms: u64,
    expires_ms: u64,
}

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

#[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(Clone, Debug, Deserialize)]
struct StreamUpsertReq {
    stream_id: String,
    title: String,
    relay_url: Option<String>,
    broadcast_name: Option<String>,
    track_name: Option<String>,
    relays: Option<Vec<PublicStreamRelay>>,
    expires_ms: Option<u64>,
}

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

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

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);
    state.streams.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();
    }
    if state.streams.len() > 1000 {
        let mut items = state.streams.values().cloned().collect::<Vec<_>>();
        items.sort_by_key(|e| std::cmp::Reverse(e.updated_ms));
        items.truncate(1000);
        state.streams = 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 }))
}

fn push_stream_relay(relays: &mut Vec<PublicStreamRelay>, relay: PublicStreamRelay) {
    if relay.relay_url.is_empty() || relay.broadcast_name.is_empty() {
        return;
    }
    if relays.iter().any(|existing| {
        existing.relay_url == relay.relay_url
            && existing.broadcast_name == relay.broadcast_name
            && existing.track_name == relay.track_name
    }) {
        return;
    }
    if relays.len() < 16 {
        relays.push(relay);
    }
}

fn normalize_stream_relays(body: &StreamUpsertReq) -> Vec<PublicStreamRelay> {
    let mut relays = Vec::new();
    if let (Some(relay_url), Some(broadcast_name)) = (&body.relay_url, &body.broadcast_name) {
        push_stream_relay(
            &mut relays,
            PublicStreamRelay {
                relay_url: clamp_str(relay_url.clone(), 512),
                broadcast_name: clamp_str(broadcast_name.clone(), 256),
                track_name: clamp_str(
                    body.track_name
                        .clone()
                        .unwrap_or_else(|| "video0.m4s".to_string()),
                    256,
                ),
            },
        );
    }
    if let Some(body_relays) = &body.relays {
        for relay in body_relays {
            push_stream_relay(
                &mut relays,
                PublicStreamRelay {
                    relay_url: clamp_str(relay.relay_url.clone(), 512),
                    broadcast_name: clamp_str(relay.broadcast_name.clone(), 256),
                    track_name: clamp_str(
                        if relay.track_name.is_empty() {
                            "video0.m4s".to_string()
                        } else {
                            relay.track_name.clone()
                        },
                        256,
                    ),
                },
            );
        }
    }
    relays
}

async fn public_streams(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.streams.values().cloned().collect::<Vec<_>>();
    entries.sort_by_key(|e| std::cmp::Reverse(e.updated_ms));
    (json_headers(), Json(PublicStreamList { now_ms: now, entries }))
}

async fn stream_upsert(
    state: axum::extract::State<Arc<RwLock<State>>>,
    Json(body): Json<StreamUpsertReq>,
) -> impl IntoResponse {
    let now = now_ms();
    let relays = normalize_stream_relays(&body);

    if body.stream_id.is_empty()
        || body.title.is_empty()
        || body.relay_url.as_deref().unwrap_or_default().is_empty()
        || body.broadcast_name.as_deref().unwrap_or_default().is_empty()
    {
        let resp =
            serde_json::json!({ "error": "missing stream_id/title/relay_url/broadcast_name" });
        return (StatusCode::BAD_REQUEST, 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 primary = relays[0].clone();

    let entry = PublicStreamEntry {
        stream_id: clamp_str(body.stream_id, 256),
        title: clamp_str(body.title, 128),
        relay_url: primary.relay_url,
        broadcast_name: primary.broadcast_name,
        track_name: primary.track_name,
        relays,
        updated_ms: now,
        expires_ms: now + ttl_ms,
    };

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

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

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/public-streams", get(public_streams))
        .route("/api/stream-upsert", post(stream_upsert))
        .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;
}
