use axum::{
    body::{Body, Bytes},
    extract::State,
    http::{
        header::{ACCEPT_RANGES, CONTENT_LENGTH, CONTENT_RANGE, RANGE},
        HeaderMap, HeaderValue, StatusCode,
    },
    response::Response,
    routing::get,
    Router,
};
use futures_util::stream;
use sha2::{Digest, Sha256};
use std::{
    convert::Infallible,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::{Duration, Instant},
};
use firelink_lib::download::{DownloadCmd, DownloadCoordinator, DownloadEvent, DownloadPayload};
use tempfile::TempDir;
use tokio::{net::TcpListener, sync::mpsc, task::JoinHandle};
use uuid::Uuid;

const TEST_TIMEOUT: Duration = Duration::from_secs(10);

#[derive(Clone)]
struct ServerState {
    data: Arc<Vec<u8>>,
    chunk_size: usize,
    chunk_delay: Duration,
    failures_remaining: Arc<AtomicUsize>,
    requests: Arc<AtomicUsize>,
    chunks_served: Arc<AtomicUsize>,
    range_starts: Arc<Mutex<Vec<Option<usize>>>>,
}

struct TestServer {
    base_url: String,
    state: ServerState,
    task: JoinHandle<()>,
}

impl TestServer {
    async fn spawn(
        data: Vec<u8>,
        chunk_size: usize,
        chunk_delay: Duration,
        failures_before_success: usize,
    ) -> Self {
        let state = ServerState {
            data: Arc::new(data),
            chunk_size,
            chunk_delay,
            failures_remaining: Arc::new(AtomicUsize::new(failures_before_success)),
            requests: Arc::new(AtomicUsize::new(0)),
            chunks_served: Arc::new(AtomicUsize::new(0)),
            range_starts: Arc::new(Mutex::new(Vec::new())),
        };
        let app = Router::new()
            .route("/file", get(serve_file))
            .route("/always-fail", get(always_fail))
            .with_state(state.clone());
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let address = listener.local_addr().unwrap();
        let task = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        Self {
            base_url: format!("http://{address}"),
            state,
            task,
        }
    }

    fn file_url(&self) -> String {
        format!("{}/file", self.base_url)
    }

    fn failure_url(&self) -> String {
        format!("{}/always-fail", self.base_url)
    }
}

impl Drop for TestServer {
    fn drop(&mut self) {
        self.task.abort();
    }
}

async fn serve_file(State(state): State<ServerState>, headers: HeaderMap) -> Response<Body> {
    state.requests.fetch_add(1, Ordering::SeqCst);
    if state
        .failures_remaining
        .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| {
            remaining.checked_sub(1)
        })
        .is_ok()
    {
        return response(StatusCode::SERVICE_UNAVAILABLE, Body::empty(), 0, None);
    }

    let requested_start = headers
        .get(RANGE)
        .and_then(|value| value.to_str().ok())
        .and_then(parse_range_start);
    state.range_starts.lock().unwrap().push(requested_start);

    let start = requested_start.unwrap_or(0);
    if start >= state.data.len() {
        return response(
            StatusCode::RANGE_NOT_SATISFIABLE,
            Body::empty(),
            0,
            Some(format!("bytes */{}", state.data.len())),
        );
    }

    let status = if requested_start.is_some() {
        StatusCode::PARTIAL_CONTENT
    } else {
        StatusCode::OK
    };
    let content_range = requested_start.map(|_| {
        format!(
            "bytes {start}-{}/{}",
            state.data.len() - 1,
            state.data.len()
        )
    });
    let content_length = state.data.len() - start;
    let data = state.data.clone();
    let chunk_size = state.chunk_size;
    let chunk_delay = state.chunk_delay;
    let chunks_served = state.chunks_served.clone();
    let body_stream = stream::unfold(start, move |offset| {
        let data = data.clone();
        let chunks_served = chunks_served.clone();
        async move {
            if offset >= data.len() {
                return None;
            }
            if !chunk_delay.is_zero() {
                tokio::time::sleep(chunk_delay).await;
            }
            let end = (offset + chunk_size).min(data.len());
            chunks_served.fetch_add(1, Ordering::SeqCst);
            let chunk = Bytes::copy_from_slice(&data[offset..end]);
            Some((Ok::<_, Infallible>(chunk), end))
        }
    });

    response(
        status,
        Body::from_stream(body_stream),
        content_length,
        content_range,
    )
}

async fn always_fail(State(state): State<ServerState>) -> Response<Body> {
    state.requests.fetch_add(1, Ordering::SeqCst);
    response(StatusCode::INTERNAL_SERVER_ERROR, Body::empty(), 0, None)
}

fn response(
    status: StatusCode,
    body: Body,
    content_length: usize,
    content_range: Option<String>,
) -> Response<Body> {
    let mut response = Response::new(body);
    *response.status_mut() = status;
    response
        .headers_mut()
        .insert(ACCEPT_RANGES, HeaderValue::from_static("bytes"));
    response.headers_mut().insert(
        CONTENT_LENGTH,
        HeaderValue::from_str(&content_length.to_string()).unwrap(),
    );
    if let Some(content_range) = content_range {
        response.headers_mut().insert(
            CONTENT_RANGE,
            HeaderValue::from_str(&content_range).unwrap(),
        );
    }
    response
}

fn parse_range_start(value: &str) -> Option<usize> {
    value
        .strip_prefix("bytes=")?
        .strip_suffix('-')?
        .parse()
        .ok()
}

fn fixture_data(size: usize) -> Vec<u8> {
    (0..size)
        .map(|index| ((index.wrapping_mul(31) + index / 7) % 251) as u8)
        .collect()
}

fn payload(id: Uuid, url: String, output_path: std::path::PathBuf) -> DownloadPayload {
    DownloadPayload {
        id,
        urls: vec![url],
        output_path,
        speed_limit: None,
        username: None,
        password: None,
        headers: None,
        cookies: None,
        user_agent: Some("Firelink integration test".to_string()),
        max_tries: 1,
        proxy: None,
    }
}

async fn next_event(events: &mut mpsc::UnboundedReceiver<DownloadEvent>) -> DownloadEvent {
    tokio::time::timeout(TEST_TIMEOUT, events.recv())
        .await
        .expect("timed out waiting for download event")
        .expect("download event channel closed")
}

async fn wait_for_progress(
    events: &mut mpsc::UnboundedReceiver<DownloadEvent>,
    id: Uuid,
    minimum_completed: u64,
) {
    loop {
        match next_event(events).await {
            DownloadEvent::Progress {
                id: event_id,
                completed,
                ..
            } if event_id == id && completed >= minimum_completed => return,
            DownloadEvent::Failed { error, .. } => panic!("download failed: {error}"),
            DownloadEvent::Completed(event_id) if event_id == id => {
                panic!("download completed before it could be paused")
            }
            _ => {}
        }
    }
}

async fn wait_for_completion(
    events: &mut mpsc::UnboundedReceiver<DownloadEvent>,
    id: Uuid,
) -> Vec<DownloadEvent> {
    let mut observed = Vec::new();
    loop {
        let event = next_event(events).await;
        match &event {
            DownloadEvent::Completed(event_id) if *event_id == id => {
                observed.push(event);
                return observed;
            }
            DownloadEvent::Failed {
                id: event_id,
                error,
            } if *event_id == id => panic!("download failed: {error}"),
            _ => observed.push(event),
        }
    }
}

fn sha256(data: &[u8]) -> Vec<u8> {
    Sha256::digest(data).to_vec()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn pause_then_resume_uses_range_and_preserves_integrity() {
    let expected = fixture_data(4 * 1024 * 1024);
    let server = TestServer::spawn(expected.clone(), 16 * 1024, Duration::from_millis(4), 0).await;
    let temp = TempDir::new().unwrap();
    let output_path = temp.path().join("paused.bin");
    let id = Uuid::from_u128(1);
    let download = payload(id, server.file_url(), output_path.clone());
    let (coordinator, mut events) = DownloadCoordinator::spawn_headless();

    coordinator
        .send(DownloadCmd::Start(download))
        .await
        .unwrap();
    wait_for_progress(&mut events, id, 256 * 1024).await;
    coordinator.send(DownloadCmd::Pause(id)).await.unwrap();
    tokio::time::sleep(Duration::from_millis(100)).await;

    let paused_len = tokio::fs::metadata(&output_path).await.unwrap().len();
    assert!(paused_len >= 256 * 1024);
    assert!(paused_len < expected.len() as u64);

    coordinator
        .send(DownloadCmd::Start(payload(
            id,
            server.file_url(),
            output_path.clone(),
        )))
        .await
        .unwrap();
    wait_for_completion(&mut events, id).await;

    let downloaded = tokio::fs::read(&output_path).await.unwrap();
    assert_eq!(downloaded.len(), expected.len());
    assert_eq!(sha256(&downloaded), sha256(&expected));
    let range_starts = server.state.range_starts.lock().unwrap().clone();
    assert_eq!(range_starts.first(), Some(&None));
    assert!(
        range_starts
            .iter()
            .skip(1)
            .flatten()
            .any(|start| *start == paused_len as usize),
        "resume request did not start at the paused file length: {range_starts:?}"
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn aggregates_many_http_chunks_with_complete_progress() {
    let expected = fixture_data(3 * 1024 * 1024);
    let server = TestServer::spawn(expected.clone(), 8 * 1024, Duration::ZERO, 0).await;
    let temp = TempDir::new().unwrap();
    let output_path = temp.path().join("chunked.bin");
    let id = Uuid::from_u128(2);
    let (coordinator, mut events) = DownloadCoordinator::spawn_headless();
    let started = Instant::now();

    coordinator
        .send(DownloadCmd::Start(payload(
            id,
            server.file_url(),
            output_path.clone(),
        )))
        .await
        .unwrap();
    let observed = wait_for_completion(&mut events, id).await;

    assert!(
        started.elapsed() < Duration::from_secs(5),
        "local 3 MiB transfer exceeded the performance budget"
    );
    assert!(server.state.chunks_served.load(Ordering::SeqCst) > 100);
    let final_progress = observed.iter().rev().find_map(|event| match event {
        DownloadEvent::Progress {
            id: event_id,
            fraction,
            completed,
            total,
        } if *event_id == id => Some((*fraction, *completed, *total)),
        _ => None,
    });
    assert_eq!(
        final_progress,
        Some((1.0, expected.len() as u64, Some(expected.len() as u64)))
    );
    assert_eq!(
        sha256(&tokio::fs::read(output_path).await.unwrap()),
        sha256(&expected)
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn retries_transient_http_failures_then_completes() {
    let expected = fixture_data(512 * 1024);
    let server = TestServer::spawn(expected.clone(), 32 * 1024, Duration::ZERO, 2).await;
    let temp = TempDir::new().unwrap();
    let output_path = temp.path().join("retry.bin");
    let id = Uuid::from_u128(3);
    let (coordinator, mut events) = DownloadCoordinator::spawn_headless();
    let mut download = payload(id, server.file_url(), output_path.clone());
    download.max_tries = 3;

    coordinator
        .send(DownloadCmd::Start(download))
        .await
        .unwrap();
    wait_for_completion(&mut events, id).await;

    assert_eq!(server.state.requests.load(Ordering::SeqCst), 3);
    assert_eq!(
        sha256(&tokio::fs::read(output_path).await.unwrap()),
        sha256(&expected)
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn reports_terminal_http_errors_after_retry_budget() {
    let server = TestServer::spawn(Vec::new(), 8 * 1024, Duration::ZERO, 0).await;
    let temp = TempDir::new().unwrap();
    let output_path = temp.path().join("failed.bin");
    let id = Uuid::from_u128(4);
    let (coordinator, mut events) = DownloadCoordinator::spawn_headless();
    let mut download = payload(id, server.failure_url(), output_path.clone());
    download.max_tries = 2;

    coordinator
        .send(DownloadCmd::Start(download))
        .await
        .unwrap();
    let error = loop {
        match next_event(&mut events).await {
            DownloadEvent::Failed {
                id: event_id,
                error,
            } if event_id == id => break error,
            DownloadEvent::Completed(event_id) if event_id == id => {
                panic!("failed download was reported as complete")
            }
            _ => {}
        }
    };

    assert!(error.contains("500 Internal Server Error"));
    assert_eq!(server.state.requests.load(Ordering::SeqCst), 2);
    assert!(!output_path.exists());
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn cancel_removes_partial_file_without_terminal_success_event() {
    let expected = fixture_data(4 * 1024 * 1024);
    let server = TestServer::spawn(expected, 16 * 1024, Duration::from_millis(4), 0).await;
    let temp = TempDir::new().unwrap();
    let output_path = temp.path().join("cancelled.bin");
    let id = Uuid::from_u128(5);
    let (coordinator, mut events) = DownloadCoordinator::spawn_headless();

    coordinator
        .send(DownloadCmd::Start(payload(
            id,
            server.file_url(),
            output_path.clone(),
        )))
        .await
        .unwrap();
    wait_for_progress(&mut events, id, 256 * 1024).await;
    coordinator.send(DownloadCmd::Cancel(id)).await.unwrap();

    tokio::time::timeout(TEST_TIMEOUT, async {
        while output_path.exists() {
            tokio::time::sleep(Duration::from_millis(10)).await;
        }
    })
    .await
    .expect("cancelled partial file was not removed");

    while let Ok(Some(event)) =
        tokio::time::timeout(Duration::from_millis(100), events.recv()).await
    {
        assert!(
            !matches!(
                event,
                DownloadEvent::Completed(event_id) if event_id == id
            ),
            "cancelled download emitted a completion event"
        );
        assert!(
            !matches!(
                event,
                DownloadEvent::Failed { id: event_id, .. } if event_id == id
            ),
            "cancelled download emitted a failure event"
        );
    }
}
