use axum::{
    body::Body,
    extract::Query,
    http::{header, StatusCode},
    response::{Html, IntoResponse, Response, Json},
    routing::get,
    Router,
};
use nokhwa::{
    pixel_format::RgbFormat,
    utils::{CameraIndex, ControlValueDescription, ControlValueSetter, KnownCameraControl, RequestedFormat, RequestedFormatType, Resolution},
    Camera,
};
use sha2::{Digest, Sha256};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use serde_json::json;

// Throughput scales with resolution - at 1080p60: ~23 MB/s conditioned, ~3 Gbps raw
// With 4K60: ~93 MB/s conditioned, ~12 Gbps raw quantum noise
const MAX_BYTES_PER_REQUEST: usize = 1024 * 1024; // 1MB per request (high-res enables this)
const CHUNK_SIZE: usize = 256; // Bytes of LSB data per hash (8:1 conditioning ratio)
const MAX_CONCURRENT: usize = 4;
const DEFAULT_PORT: u16 = 8787;

static ACTIVE_REQUESTS: AtomicUsize = AtomicUsize::new(0);
static REQUEST_COUNTER: AtomicU64 = AtomicU64::new(0);

fn is_fake_camera() -> bool {
    std::env::var("FAKE_CAMERA")
        .map(|v| v == "1" || v.to_lowercase() == "true")
        .unwrap_or(false)
}

/// Get requested resolution from environment, defaulting to 1080p for high throughput
fn get_resolution() -> (u32, u32) {
    let width = std::env::var("CAMERA_WIDTH")
        .ok()
        .and_then(|w| w.parse().ok())
        .unwrap_or(1920);
    let height = std::env::var("CAMERA_HEIGHT")
        .ok()
        .and_then(|h| h.parse().ok())
        .unwrap_or(1080);
    (width, height)
}

#[derive(serde::Deserialize)]
struct RandomQuery {
    #[serde(default = "default_bytes")]
    bytes: usize,
    #[serde(default)]
    hex: bool,
}

fn default_bytes() -> usize { 32 }

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let port = std::env::var("PORT").ok().and_then(|p| p.parse().ok()).unwrap_or(DEFAULT_PORT);
    let (width, height) = get_resolution();
    
    if is_fake_camera() {
        println!("FAKE_CAMERA mode enabled - using /dev/urandom for entropy");
    } else {
        println!("Testing camera access...");
        match test_camera(width, height) {
            Ok((actual_w, actual_h, frame_size)) => {
                let conditioned_per_frame = (frame_size / CHUNK_SIZE) * 32;
                let throughput_30fps = conditioned_per_frame * 30;
                let raw_gbps = (frame_size as f64 * 30.0 * 8.0) / 1_000_000_000.0;
                println!("Camera OK at {}x{} - {} bytes/frame", actual_w, actual_h, frame_size);
                println!("Raw throughput: {:.1} Gbps at 30fps", raw_gbps);
                println!("Conditioned output: ~{} MB/s at 30fps (8:1 ratio)", throughput_30fps / 1_000_000);
                println!("Ensure lens is covered for optimal quantum noise capture");
            }
            Err(e) => {
                eprintln!("Camera error: {}. Server will still start.", e);
            }
        }
    }

    let app = Router::new()
        .route("/", get(index))
        .route("/random", get(get_random))
        .route("/health", get(health))
        .route("/.well-known/mcp.json", get(mcp_wellknown));

    let addr = format!("0.0.0.0:{}", port);
    println!("Camera QRNG (LavaRnd-style) on http://{}", addr);
    let listener = tokio::net::TcpListener::bind(&addr).await?;
    axum::serve(listener, app).await?;
    Ok(())
}

async fn index() -> Html<&'static str> { Html(INDEX_HTML) }
async fn health() -> &'static str { "ok" }

async fn mcp_wellknown() -> Json<serde_json::Value> {
    Json(json!({
        "mcp": {
            "spec_version": "2026-01-21",
            "status": "active",
            "servers": [],
            "tools": [{
                "name": "camera-qrng",
                "description": "High-throughput quantum RNG using thermal noise from covered camera sensor - Gbps of raw quantum entropy",
                "url_template": "{origin}/random?bytes={bytes}&hex={hex}",
                "capabilities": ["random-generation", "entropy-source", "quantum"],
                "auth": { "type": "none" },
                "parameters": {
                    "bytes": { "type": "integer", "default": 32, "max": 1048576, "description": "Number of random bytes (up to 1MB)" },
                    "hex": { "type": "boolean", "default": false, "description": "Return hex-encoded string" }
                }
            }]
        }
    }))
}

/// Test camera and return (width, height, frame_size) on success
fn test_camera(req_width: u32, req_height: u32) -> Result<(u32, u32, usize), String> {
    let index = CameraIndex::Index(0);
    let resolution = Resolution::new(req_width, req_height);
    let format = RequestedFormat::new::<RgbFormat>(RequestedFormatType::HighestResolution(resolution));
    let mut camera = Camera::new(index, format).map_err(|e| e.to_string())?;
    camera.open_stream().map_err(|e| e.to_string())?;
    let frame = camera.frame().map_err(|e| e.to_string())?;
    let res = camera.resolution();
    let frame_size = frame.buffer().len();
    camera.stop_stream().ok();
    Ok((res.width(), res.height(), frame_size))
}

/// Extract maximum value from a ControlValueDescription if it's an integer range
fn get_max_int(desc: &ControlValueDescription) -> Option<i64> {
    match desc {
        ControlValueDescription::IntegerRange { max, .. } => Some(*max),
        ControlValueDescription::Integer { value, .. } => Some(*value),
        _ => None,
    }
}

/// Configure camera for optimal quantum noise capture (LavaRnd approach).
/// Maximizes gain and brightness to amplify dark current and thermal noise.
fn configure_for_thermal_noise(camera: &mut Camera) {
    // Maximize gain to amplify thermal/quantum noise
    if let Ok(ctrl) = camera.camera_control(KnownCameraControl::Gain) {
        if let Some(max) = get_max_int(ctrl.description()) {
            let _ = camera.set_camera_control(
                KnownCameraControl::Gain,
                ControlValueSetter::Integer(max),
            );
        }
    }
    
    // Maximize brightness
    if let Ok(ctrl) = camera.camera_control(KnownCameraControl::Brightness) {
        if let Some(max) = get_max_int(ctrl.description()) {
            let _ = camera.set_camera_control(
                KnownCameraControl::Brightness,
                ControlValueSetter::Integer(max),
            );
        }
    }
    
    // Set exposure to maximum if available (longer exposure = more thermal noise accumulation)
    if let Ok(ctrl) = camera.camera_control(KnownCameraControl::Exposure) {
        if let Some(max) = get_max_int(ctrl.description()) {
            let _ = camera.set_camera_control(
                KnownCameraControl::Exposure,
                ControlValueSetter::Integer(max),
            );
        }
    }
}

async fn get_random(Query(params): Query<RandomQuery>) -> Response {
    let current = ACTIVE_REQUESTS.fetch_add(1, Ordering::SeqCst);
    if current >= MAX_CONCURRENT {
        ACTIVE_REQUESTS.fetch_sub(1, Ordering::SeqCst);
        return (StatusCode::TOO_MANY_REQUESTS, "Too many requests").into_response();
    }
    let bytes = params.bytes.min(MAX_BYTES_PER_REQUEST);
    if bytes == 0 {
        ACTIVE_REQUESTS.fetch_sub(1, Ordering::SeqCst);
        return (StatusCode::BAD_REQUEST, "bytes must be > 0").into_response();
    }
    let request_id = REQUEST_COUNTER.fetch_add(1, Ordering::SeqCst);
    let use_fake = is_fake_camera();
    
    let result = tokio::task::spawn_blocking(move || {
        if use_fake {
            extract_entropy_fake(bytes, request_id)
        } else {
            extract_entropy_camera(bytes, request_id)
        }
    }).await;
    ACTIVE_REQUESTS.fetch_sub(1, Ordering::SeqCst);

    match result {
        Ok(Ok(data)) => {
            if params.hex {
                Response::builder().header(header::CONTENT_TYPE, "text/plain")
                    .body(Body::from(hex::encode(&data))).unwrap()
            } else {
                Response::builder().header(header::CONTENT_TYPE, "application/octet-stream")
                    .body(Body::from(data)).unwrap()
            }
        }
        Ok(Err(e)) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(),
    }
}

/// Fake entropy source using /dev/urandom - for testing without camera hardware.
/// Simulates high-resolution camera frames for realistic throughput testing.
fn extract_entropy_fake(num_bytes: usize, request_id: u64) -> Result<Vec<u8>, String> {
    use std::io::Read;
    
    let (width, height) = get_resolution();
    let frame_size = (width * height * 3) as usize;
    
    let mut entropy = Vec::with_capacity(num_bytes);
    let mut hasher = Sha256::new();
    
    let mut urandom = std::fs::File::open("/dev/urandom").map_err(|e| e.to_string())?;
    let mut fake_frame = vec![0u8; frame_size];
    
    let mut frame_idx: u64 = 0;
    while entropy.len() < num_bytes {
        urandom.read_exact(&mut fake_frame).map_err(|e| e.to_string())?;
        
        // Extract LSBs (2 bits per byte - highest entropy density)
        let lsbs: Vec<u8> = fake_frame.iter().map(|b| b & 0x03).collect();
        
        // Hash in chunks - each CHUNK_SIZE bytes of LSBs produces 32 bytes output
        for (chunk_idx, chunk) in lsbs.chunks(CHUNK_SIZE).enumerate() {
            hasher.update(chunk);
            hasher.update(&request_id.to_le_bytes());
            hasher.update(&frame_idx.to_le_bytes());
            hasher.update(&(chunk_idx as u64).to_le_bytes());
            hasher.update(&nanos_now().to_le_bytes());
            
            entropy.extend_from_slice(&hasher.finalize_reset());
            
            if entropy.len() >= num_bytes {
                break;
            }
        }
        frame_idx += 1;
    }
    
    entropy.truncate(num_bytes);
    Ok(entropy)
}

/// Extract entropy from camera quantum noise using chunked SHA-256 conditioning.
/// 
/// Throughput scales with camera resolution:
/// - 640x480 @ 30fps:  ~27 MB/s raw (~216 Mbps), ~3.4 MB/s conditioned
/// - 1080p @ 30fps:    ~186 MB/s raw (~1.5 Gbps), ~23 MB/s conditioned
/// - 1080p @ 60fps:    ~373 MB/s raw (~3 Gbps), ~47 MB/s conditioned
/// - 4K @ 30fps:       ~746 MB/s raw (~6 Gbps), ~93 MB/s conditioned
/// - 4K @ 60fps:       ~1.49 GB/s raw (~12 Gbps), ~186 MB/s conditioned
fn extract_entropy_camera(num_bytes: usize, request_id: u64) -> Result<Vec<u8>, String> {
    let (req_width, req_height) = get_resolution();
    let index = CameraIndex::Index(0);
    let resolution = Resolution::new(req_width, req_height);
    let format = RequestedFormat::new::<RgbFormat>(RequestedFormatType::HighestResolution(resolution));
    let mut camera = Camera::new(index, format).map_err(|e| e.to_string())?;
    camera.open_stream().map_err(|e| e.to_string())?;
    
    // Configure camera for quantum noise capture (high gain, max brightness)
    configure_for_thermal_noise(&mut camera);
    
    let mut entropy = Vec::with_capacity(num_bytes);
    let mut hasher = Sha256::new();
    
    let mut frame_idx: u64 = 0;
    while entropy.len() < num_bytes {
        let frame = camera.frame().map_err(|e| e.to_string())?;
        let raw = frame.buffer();
        
        // Extract LSBs (2 bits per byte - highest entropy density in quantum noise)
        let lsbs: Vec<u8> = raw.iter().map(|b| b & 0x03).collect();
        
        // Hash in chunks - each CHUNK_SIZE bytes produces 32 bytes conditioned output
        // At 1080p: ~24,300 chunks/frame = ~778 KB conditioned per frame
        // At 4K: ~97,200 chunks/frame = ~3.1 MB conditioned per frame
        for (chunk_idx, chunk) in lsbs.chunks(CHUNK_SIZE).enumerate() {
            hasher.update(chunk);
            hasher.update(&request_id.to_le_bytes());
            hasher.update(&frame_idx.to_le_bytes());
            hasher.update(&(chunk_idx as u64).to_le_bytes());
            hasher.update(&nanos_now().to_le_bytes());
            
            entropy.extend_from_slice(&hasher.finalize_reset());
            
            if entropy.len() >= num_bytes {
                break;
            }
        }
        frame_idx += 1;
    }
    
    camera.stop_stream().ok();
    entropy.truncate(num_bytes);
    Ok(entropy)
}

fn nanos_now() -> u128 {
    std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos()
}

const INDEX_HTML: &str = include_str!("index.html");