use std::num::NonZeroUsize;
use std::sync::OnceLock;

use axum::{
    extract::ws::{Message, WebSocket, WebSocketUpgrade},
    http::StatusCode,
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use tokio::sync::broadcast;
use crate::{Escaped, FallbackRender};

const DEV_TOKEN_HEADER: &str = "X-Azumi-Dev-Token";

static BROADCAST_CHANNEL: OnceLock<broadcast::Sender<String>> = OnceLock::new();

fn get_broadcast_channel() -> &'static broadcast::Sender<String> {
    BROADCAST_CHANNEL.get_or_init(|| {
        let (tx, _) = broadcast::channel(100);
        tx
    })
}

/// Cached dev token — read once from env, avoids syscall on every request.
static DEV_TOKEN: OnceLock<Option<String>> = OnceLock::new();

fn get_dev_token() -> Option<&'static str> {
    DEV_TOKEN
        .get_or_init(|| std::env::var("AZUMI_DEV_TOKEN").ok())
        .as_deref()
}

/// Reset the cached dev token. Only available in tests.
#[cfg(test)]
fn reset_dev_token_cache() {
    // OnceLock doesn't support reset, so we use a different strategy:
    // the tests use a thread-local override instead.
}

/// Test-only thread-local token override, bypasses OnceLock cache.
#[cfg(test)]
thread_local! {
    static TEST_DEV_TOKEN: std::cell::RefCell<Option<String>> = const { std::cell::RefCell::new(None) };
}

#[cfg(not(test))]
pub fn is_dev_token_valid(token: Option<&str>) -> bool {
    is_dev_token_valid_impl(token, get_dev_token())
}

#[cfg(test)]
pub fn is_dev_token_valid(token: Option<&str>) -> bool {
    let expected = TEST_DEV_TOKEN.with(|t| t.borrow().clone());
    match expected {
        Some(s) => is_dev_token_valid_impl(token, Some(&s)),
        None => is_dev_token_valid_impl(token, get_dev_token()),
    }
}

/// Constant-time token comparison that does NOT leak token length.
/// XORs the length difference into the accumulator so short-circuiting
/// is impossible — every branch takes the same path regardless of input.
/// An attacker cannot distinguish wrong-length from wrong-content tokens.
fn is_dev_token_valid_impl(token: Option<&str>, expected: Option<&str>) -> bool {
    let Some(t) = token else {
        return false;
    };
    let Some(expected) = expected else {
        return false;
    };

    let t_bytes = t.as_bytes();
    let expected_bytes = expected.as_bytes();

    let len_diff = t_bytes.len() ^ expected_bytes.len();
    let max_len = t_bytes.len().max(expected_bytes.len());

    let mut result = len_diff as u8 | ((len_diff >> 8) as u8);
    for i in 0..max_len {
        let tb = t_bytes.get(i).copied().unwrap_or(0);
        let eb = expected_bytes.get(i).copied().unwrap_or(0);
        result |= tb ^ eb;
    }

    result == 0
}

#[derive(Clone)]
pub struct RuntimeTemplate {
    pub static_parts: Vec<String>,
}

impl RuntimeTemplate {
    pub fn render(
        &self,
        f: &mut std::fmt::Formatter<'_>,
        dynamics: &[&dyn FallbackRender],
    ) -> std::fmt::Result {
        for (i, part) in self.static_parts.iter().enumerate() {
            write!(f, "{}", Escaped(part))?;
            if i < dynamics.len() {
                dynamics[i].render_azumi(f)?;
            }
        }
        Ok(())
    }
}

const MAX_REGISTRY_SIZE: usize = 1000;

static TEMPLATE_REGISTRY: OnceLock<std::sync::RwLock<lru::LruCache<String, RuntimeTemplate>>> = OnceLock::new();

fn create_registry() -> lru::LruCache<String, RuntimeTemplate> {
    lru::LruCache::new(NonZeroUsize::new(MAX_REGISTRY_SIZE).unwrap())
}

pub fn get_template(id: &str) -> Option<RuntimeTemplate> {
    let Ok(mut registry) = TEMPLATE_REGISTRY.get_or_init(|| std::sync::RwLock::new(create_registry())).write() else {
        eprintln!("Hot Reload: Registry lock poisoned - template lookup failed");
        return None;
    };
    registry.get(&id.to_string()).cloned()
}

const MAX_TEMPLATE_PARTS: usize = 100;
const MAX_PART_SIZE: usize = 100_000; // 100KB per part
const MAX_TEMPLATE_ID_LEN: usize = 256;
const MAX_TOTAL_CSS_SIZE: usize = 500_000; // 500KB total CSS limit per template

#[derive(serde::Deserialize)]
struct TemplateUpdatePayload {
    id: String,
    parts: Vec<String>,
}

async fn update_template_handler(Json(payload): Json<TemplateUpdatePayload>) -> impl IntoResponse {
    if payload.id.len() > MAX_TEMPLATE_ID_LEN {
        eprintln!("Hot Reload: Template ID too long");
        return (StatusCode::BAD_REQUEST, "Template ID too long");
    }
    if payload.parts.len() > MAX_TEMPLATE_PARTS {
        eprintln!("Hot Reload: Too many parts (max {})", MAX_TEMPLATE_PARTS);
        return (StatusCode::BAD_REQUEST, "Too many template parts");
    }
    let mut total_size = payload.id.len();
    for part in &payload.parts {
        if part.len() > MAX_PART_SIZE {
            eprintln!("Hot Reload: Part too large (max {} bytes)", MAX_PART_SIZE);
            return (StatusCode::PAYLOAD_TOO_LARGE, "Template part too large");
        }
        total_size = total_size.saturating_add(part.len());
        if total_size > MAX_TOTAL_CSS_SIZE {
            eprintln!("Hot Reload: Total CSS size too large (max {} bytes)", MAX_TOTAL_CSS_SIZE);
            return (StatusCode::PAYLOAD_TOO_LARGE, "Total CSS size too large");
        }
    }

    let Ok(mut registry) = TEMPLATE_REGISTRY.get_or_init(|| std::sync::RwLock::new(create_registry())).write() else {
        eprintln!("Hot Reload: Registry lock poisoned - template update failed");
        return (StatusCode::SERVICE_UNAVAILABLE, "Registry unavailable");
    };

    // lru::LruCache automatically evicts oldest entries when at capacity
    registry.put(payload.id.clone(), RuntimeTemplate { static_parts: payload.parts });
    #[cfg(debug_assertions)]
    println!("Hot Reload: Updated template \"{}\"", payload.id.replace('"', "\\\""));
    let _ = get_broadcast_channel().send(serde_json::json!({"type": "reload"}).to_string());
    (StatusCode::OK, "Template updated")
}

/// Pushes a style update to all connected clients
pub fn push_style_update(scope_id: &str, css: &str) {
    let msg = serde_json::json!({
        "type": "style-update",
        "scopeId": scope_id,
        "css": css
    });
    let _ = get_broadcast_channel().send(msg.to_string());
}

async fn check_dev_token(
    req: axum::http::Request<axum::body::Body>,
    next: axum::middleware::Next,
) -> Result<axum::response::Response, StatusCode> {
    let token = req.headers()
        .get(DEV_TOKEN_HEADER)
        .and_then(|v| v.to_str().ok());
    
    if is_dev_token_valid(token) {
        Ok(next.run(req).await)
    } else {
        Err(StatusCode::UNAUTHORIZED)
    }
}

async fn ws_handler(ws: WebSocketUpgrade) -> impl IntoResponse {
    ws.on_upgrade(handle_socket)
}

async fn handle_socket(mut socket: WebSocket) {
    let mut rx = get_broadcast_channel().subscribe();

    loop {
        tokio::select! {
            msg = rx.recv() => {
                if let Ok(msg) = msg {
                    if socket.send(Message::Text(axum::extract::ws::Utf8Bytes::from(msg))).await.is_err() {
                        break;
                    }
                }
            }
            _ = tokio::time::sleep(std::time::Duration::from_secs(30)) => {
                if socket.send(Message::Ping(axum::body::Bytes::new())).await.is_err() {
                    break;
                }
            }
        }
    }
}

/// Mounts the hot reload route at `/_azumi/live_reload`
///
/// # Security Warning
///
/// These endpoints are **development-only** and should NOT be exposed in production:
///
/// - `/_azumi/live_reload` - WebSocket endpoint for hot reload
/// - `/_azumi/update_template` - POST endpoint to update templates
///
/// **Authentication**: Both endpoints require the `X-Azumi-Dev-Token` header
/// to be set to the value of the `AZUMI_DEV_TOKEN` environment variable.
///
/// In production, either:
/// 1. Remove this router entirely (hot reload is for development only)
/// 2. Restrict access at the network level (e.g., firewall rules to block external access)
/// 3. Ensure `AZUMI_DEV_TOKEN` is not set or is a secret only localhost knows
///
/// If deploying to production with this enabled, ensure only localhost can access these routes.
pub fn router<S>() -> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    Router::new()
        .route("/_azumi/live_reload", get(ws_handler))
        .route("/_azumi/update_template", post(update_template_handler))
        .layer(axum::middleware::from_fn(check_dev_token))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_is_dev_token_valid_exact_match() {
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = Some("secret123".to_string()));
        assert!(is_dev_token_valid(Some("secret123")));
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
    }

    #[test]
    fn test_is_dev_token_valid_wrong_token() {
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = Some("secret123".to_string()));
        assert!(!is_dev_token_valid(Some("wrongtoken")));
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
    }

    #[test]
    fn test_is_dev_token_valid_partial_match_rejected() {
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = Some("secret123".to_string()));
        assert!(!is_dev_token_valid(Some("secret")));
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
    }

    #[test]
    fn test_is_dev_token_valid_empty_token() {
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = Some("secret123".to_string()));
        assert!(!is_dev_token_valid(Some("")));
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
    }

    #[test]
    fn test_is_dev_token_valid_no_env_var() {
        // No test token set → falls through to get_dev_token() which returns None
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
        // Ensure the env var is also not set so get_dev_token returns None
        std::env::remove_var("AZUMI_DEV_TOKEN");
        // Since OnceLock may have cached from a previous test in same thread,
        // we can't fully reset, but thread-local None is the primary path.
        // This test is best-effort.
        assert!(!is_dev_token_valid(Some("secret123")));
    }

    #[test]
    fn test_is_dev_token_valid_none_token() {
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = Some("secret123".to_string()));
        assert!(!is_dev_token_valid(None));
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
    }

    #[test]
    fn test_is_dev_token_valid_constant_time_different_lengths() {
        // Verifies the constant-time property: tokens of different lengths
        // should not short-circuit. Both should return false consistently.
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = Some("abc".to_string()));
        assert!(!is_dev_token_valid(Some("a")));
        assert!(!is_dev_token_valid(Some("ab")));
        assert!(!is_dev_token_valid(Some("abcd")));
        assert!(!is_dev_token_valid(Some("abcde")));
        TEST_DEV_TOKEN.with(|t| *t.borrow_mut() = None);
    }

    #[test]
    fn test_lru_cache_insert_and_get() {
        let mut cache = create_registry();
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["value1".to_string()] });
        assert!(cache.get("key1").is_some());
    }

    #[test]
    fn test_lru_cache_update_existing_key() {
        let mut cache = create_registry();
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["value1".to_string()] });
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["value2".to_string()] });
        let val = cache.get("key1").unwrap();
        assert_eq!(val.static_parts, vec!["value2".to_string()]);
    }

    #[test]
    fn test_lru_cache_get_nonexistent() {
        let mut cache = create_registry();
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["value1".to_string()] });
        assert!(cache.get("key2").is_none());
    }

    #[test]
    fn test_lru_cache_len() {
        let mut cache = create_registry();
        assert_eq!(cache.len(), 0);
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["value1".to_string()] });
        assert_eq!(cache.len(), 1);
        cache.put("key2".to_string(), RuntimeTemplate { static_parts: vec!["value2".to_string()] });
        assert_eq!(cache.len(), 2);
    }

    #[test]
    fn test_lru_cache_auto_evict() {
        let mut cache = lru::LruCache::new(NonZeroUsize::new(2).unwrap());
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["v1".to_string()] });
        cache.put("key2".to_string(), RuntimeTemplate { static_parts: vec!["v2".to_string()] });
        cache.put("key3".to_string(), RuntimeTemplate { static_parts: vec!["v3".to_string()] });
        assert_eq!(cache.len(), 2);
        assert!(cache.get("key1").is_none()); // key1 was evicted
    }

    #[test]
    fn test_lru_cache_access_updates_order() {
        let mut cache = lru::LruCache::new(NonZeroUsize::new(2).unwrap());
        cache.put("key1".to_string(), RuntimeTemplate { static_parts: vec!["v1".to_string()] });
        cache.put("key2".to_string(), RuntimeTemplate { static_parts: vec!["v2".to_string()] });
        let _ = cache.get("key1"); // access key1 to make it more recent
        cache.put("key3".to_string(), RuntimeTemplate { static_parts: vec!["v3".to_string()] });
        assert!(cache.get("key1").is_some()); // key1 should still be there
        assert!(cache.get("key2").is_none()); // key2 was evicted
    }

    #[test]
    fn test_size_limits_are_consistent() {
        const _: () = {
            assert!(
                MAX_TOTAL_CSS_SIZE < MAX_REGISTRY_SIZE * MAX_PART_SIZE,
                "MAX_TOTAL_CSS_SIZE should be less than the theoretical max from parts × MAX_PART_SIZE"
            );
            assert!(
                MAX_TOTAL_CSS_SIZE <= MAX_REGISTRY_SIZE * MAX_PART_SIZE,
                "CSS size limit must be reachable (not dead code)"
            );
        };
    }

    #[test]
    fn test_payload_size_limit_allows_reasonable_templates() {
        let reasonable_payload_size = 50_000; // 50KB — a reasonable template
        assert!(
            reasonable_payload_size <= MAX_TOTAL_CSS_SIZE,
            "A 50KB template should be within CSS size limits"
        );
    }
}