//! HTTP/2 connection handling using h2

use super::handler::HttpConnectionHandler;
use crate::config::ServerConfig;
use crate::errors::{FalcornError, Result};
use crate::http::{Body, HttpMethod, HttpRequest, HttpResponse};
use crate::logging::Logger;
use crate::server::request_handler::RequestHandler;
use bytes::Bytes;
use futures::future::poll_fn;
use http::{HeaderName, HeaderValue, Request, Response};
use std::io::Read;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::task::{Context, Poll};
use tokio::io::ReadBuf;
use tokio::io::{AsyncRead, AsyncWrite};

/// HTTP/2 handler (async-only)
#[derive(Default, Clone)]
pub struct Http2Handler;

impl Http2Handler {
    pub fn new() -> Self {
        Self
    }
}

// Serve HTTP/2 requests asynchronously.
pub async fn serve_connection_async<S: AsyncRead + AsyncWrite + Unpin>(
    io: S,
    peer_addr: SocketAddr,
    handler: RequestHandler,
    logger: Logger,
    config: Arc<ServerConfig>,
    initial: Vec<u8>,
    request_count: Option<Arc<AtomicUsize>>,
) -> Result<()> {
    let mut builder = h2::server::Builder::new();

    // Conservative defaults; can be surfaced in config later.
    let window = (config.buffer_size as u32)
        .max(64 * 1024)
        .min(4 * 1024 * 1024);
    builder.initial_window_size(window);
    builder.initial_connection_window_size(window * 2);
    builder.max_concurrent_streams(config.max_connections.max(1) as u32);

    let mut connection = builder
        .handshake(PrefixedIo::new(io, initial))
        .await
        .map_err(|e| FalcornError::runtime(format!("h2 handshake error: {}", e)))?;

    while let Some(result) = connection.accept().await {
        let (req, respond) =
            result.map_err(|e| FalcornError::runtime(format!("h2 accept error: {}", e)))?;

        let handler = handler.clone();
        let logger = logger.clone();
        let request_count = request_count.clone();

        tokio::spawn(async move {
            if let Some(counter) = request_count.as_ref() {
                counter.fetch_add(1, Ordering::Relaxed);
            }

            match handle_h2_request(req, peer_addr, handler.clone()).await {
                Ok(resp) => {
                    let _ = send_h2_response(respond, resp).await;
                }
                Err(err) => {
                    logger.error(|| format!("HTTP/2 request error: {}", err));
                    let fallback = handler.handle_error_with_request(err, None);
                    let _ = send_h2_response(respond, fallback).await;
                }
            }
        });
    }

    Ok(())
}

async fn handle_h2_request(
    req: Request<h2::RecvStream>,
    peer_addr: SocketAddr,
    handler: RequestHandler,
) -> Result<HttpResponse> {
    let (parts, mut body) = req.into_parts();

    validate_h2_request_parts(&parts)?;

    let method = HttpMethod::from_str(parts.method.as_str())?;
    let path = Bytes::copy_from_slice(parts.uri.path().as_bytes());
    let query = parts
        .uri
        .query()
        .map(|q| Bytes::copy_from_slice(q.as_bytes()));

    let mut headers: Vec<(Bytes, Bytes)> = Vec::new();
    for (name, value) in parts.headers.iter() {
        if let Some((n, v)) = filter_h2_header(name, value) {
            headers.push((
                Bytes::copy_from_slice(n.as_str().as_bytes()),
                Bytes::copy_from_slice(v.as_bytes()),
            ));
        }
    }

    let mut body_bytes = Vec::new();
    while let Some(chunk) = body.data().await {
        let bytes = chunk.map_err(|e| FalcornError::runtime(format!("h2 body error: {}", e)))?;

        // Release flow-control capacity as we consume.
        body.flow_control()
            .release_capacity(bytes.len())
            .map_err(|e| FalcornError::runtime(format!("h2 flow control error: {}", e)))?;
        body_bytes.extend_from_slice(&bytes);
    }

    let request = HttpRequest {
        raw: Bytes::new(),
        method,
        path,
        query_string: query,
        version: http::Version::HTTP_2,
        headers,
        body: if body_bytes.is_empty() {
            None
        } else {
            Some(std::sync::Arc::new(Bytes::from(body_bytes)))
        },
        peer_addr,
    };

    match handler.handle_async(&request).await {
        Ok(resp) => Ok(resp),
        Err(err) => Ok(handler.handle_error_with_request(err, Some(&request))),
    }
}

async fn send_h2_response(
    mut respond: h2::server::SendResponse<Bytes>,
    mut response: HttpResponse,
) -> Result<()> {
    let mut builder = Response::builder().status(response.status);

    {
        let headers = builder
            .headers_mut()
            .ok_or_else(|| FalcornError::runtime("Failed to access response headers"))?;
        for (name, value) in response.headers.iter() {
            if let Some((n, v)) = filter_h2_header(name, value) {
                headers.insert(n, v);
            }
        }
    }

    let resp = builder
        .body(())
        .map_err(|e| FalcornError::runtime(format!("h2 response build error: {}", e)))?;

    let mut send_stream = respond
        .send_response(resp, false)
        .map_err(|e| FalcornError::runtime(format!("h2 send_response error: {}", e)))?;

    match &mut response.body {
        Body::Empty => {
            send_stream
                .send_data(Bytes::new(), true)
                .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
        }
        Body::Bytes(bytes) => {
            send_body_with_flow(&mut send_stream, bytes.clone()).await?;
        }
        Body::AsyncStream(stream) => {
            send_stream_with_flow(&mut send_stream, stream).await?;
        }
        Body::SyncReader(reader) => {
            let mut buf = Vec::new();
            reader
                .read_to_end(&mut buf)
                .map_err(FalcornError::IoError)?;
            let bytes = Bytes::from(buf);
            send_body_with_flow(&mut send_stream, bytes).await?;
        }
        Body::File { path, offset, len } => {
            send_file_with_flow(&mut send_stream, path, *offset, *len).await?;
        }
    }

    Ok(())
}

async fn send_body_with_flow(send_stream: &mut h2::SendStream<Bytes>, body: Bytes) -> Result<()> {
    if body.is_empty() {
        send_stream
            .send_data(Bytes::new(), true)
            .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
        return Ok(());
    }

    let mut offset = 0usize;
    let total = body.len();

    while offset < total {
        let cap = poll_fn(|cx| send_stream.poll_capacity(cx))
            .await
            .ok_or_else(|| FalcornError::runtime("h2 stream closed"))?;

        let cap = cap.map_err(|e| FalcornError::runtime(format!("h2 capacity error: {}", e)))?;

        if cap == 0 {
            continue;
        }

        let chunk_len = (cap as usize).min(total - offset);
        let chunk = body.slice(offset..offset + chunk_len);
        offset += chunk_len;

        let end = offset == total;
        send_stream
            .send_data(chunk, end)
            .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
    }

    Ok(())
}

async fn send_stream_with_flow(
    send_stream: &mut h2::SendStream<Bytes>,
    body: &mut crate::http::body::AsyncBodyStream,
) -> Result<()> {
    use futures::stream::StreamExt;
    while let Some(chunk) = body.next().await {
        let bytes =
            chunk.map_err(|e| FalcornError::runtime(format!("h2 body stream error: {}", e)))?;
        if bytes.is_empty() {
            continue;
        }

        let mut offset = 0usize;
        let total = bytes.len();
        while offset < total {
            let cap = poll_fn(|cx| send_stream.poll_capacity(cx))
                .await
                .ok_or_else(|| FalcornError::runtime("h2 stream closed"))?;

            let cap =
                cap.map_err(|e| FalcornError::runtime(format!("h2 capacity error: {}", e)))?;
            if cap == 0 {
                continue;
            }

            let chunk_len = (cap as usize).min(total - offset);
            let chunk = bytes.slice(offset..offset + chunk_len);
            offset += chunk_len;
            let end = false;
            send_stream
                .send_data(chunk, end)
                .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
        }
    }

    send_stream
        .send_data(Bytes::new(), true)
        .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
    Ok(())
}

async fn send_file_with_flow(
    send_stream: &mut h2::SendStream<Bytes>,
    path: &std::path::Path,
    offset: u64,
    _len: Option<u64>,
) -> Result<()> {
    use tokio::io::{AsyncReadExt, AsyncSeekExt};

    let mut file = tokio::fs::File::open(path)
        .await
        .map_err(FalcornError::IoError)?;
    if offset > 0 {
        file.seek(std::io::SeekFrom::Start(offset))
            .await
            .map_err(FalcornError::IoError)?;
    }

    let mut buf = [0u8; 16384];
    loop {
        let n = file.read(&mut buf).await.map_err(FalcornError::IoError)?;
        if n == 0 {
            break;
        }
        let bytes = Bytes::copy_from_slice(&buf[..n]);

        let mut sent = 0usize;
        while sent < n {
            let cap = poll_fn(|cx| send_stream.poll_capacity(cx))
                .await
                .ok_or_else(|| FalcornError::runtime("h2 stream closed"))?;
            let cap =
                cap.map_err(|e| FalcornError::runtime(format!("h2 capacity error: {}", e)))?;
            if cap == 0 {
                continue;
            }
            let chunk_len = (cap as usize).min(n - sent);
            let chunk = bytes.slice(sent..sent + chunk_len);
            sent += chunk_len;
            let end = false;
            send_stream
                .send_data(chunk, end)
                .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
        }
    }

    send_stream
        .send_data(Bytes::new(), true)
        .map_err(|e| FalcornError::runtime(format!("h2 send_data error: {}", e)))?;
    Ok(())
}

pub(crate) fn filter_h2_header(
    name: &HeaderName,
    value: &HeaderValue,
) -> Option<(HeaderName, HeaderValue)> {
    let n = name.as_str();

    // Forbidden connection-specific headers in HTTP/2
    if matches!(
        n,
        "connection" | "keep-alive" | "proxy-connection" | "transfer-encoding" | "upgrade"
    ) {
        return None;
    }

    // TE is only allowed with "trailers"
    if n == "te" {
        if let Ok(v) = value.to_str() {
            if v.trim().eq_ignore_ascii_case("trailers") {
                return Some((name.clone(), value.clone()));
            }
        }
        return None;
    }

    Some((name.clone(), value.clone()))
}

pub(crate) fn validate_h2_request_parts(parts: &http::request::Parts) -> Result<()> {
    if parts.method.as_str().is_empty() {
        return Err(FalcornError::http_parse(
            "HTTP/2 missing required pseudo-header :method",
        ));
    }

    if parts.uri.path().is_empty() {
        return Err(FalcornError::http_parse(
            "HTTP/2 missing required pseudo-header :path",
        ));
    }

    let has_authority = parts.uri.authority().is_some();
    let has_host = parts.headers.get(http::header::HOST).is_some();
    if !has_authority && !has_host {
        return Err(FalcornError::http_parse("HTTP/2 missing :authority/host"));
    }

    Ok(())
}

impl HttpConnectionHandler for Http2Handler {
    fn handle_sync<S: std::io::Read + std::io::Write + super::http1::ReadTimeout>(
        &self,
        _stream: S,
        _peer_addr: SocketAddr,
        _handler: &RequestHandler,
        _logger: &Logger,
        _config: &ServerConfig,
        _initial: Vec<u8>,
        _request_count: Option<&AtomicUsize>,
    ) -> Result<()> {
        Err(FalcornError::runtime(
            "HTTP/2 requires async worker (tokio). Use worker_class=async",
        ))
    }

    fn handle_async<'a, S>(
        &'a self,
        stream: S,
        peer_addr: SocketAddr,
        handler: RequestHandler,
        logger: Logger,
        config: Arc<ServerConfig>,
        initial: Vec<u8>,
        request_count: Option<Arc<AtomicUsize>>,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + 'a>>
    where
        S: AsyncRead + AsyncWrite + Unpin + Send + 'a,
    {
        Box::pin(serve_connection_async(
            stream,
            peer_addr,
            handler,
            logger,
            config,
            initial,
            request_count,
        ))
    }
}

struct PrefixedIo<S> {
    inner: S,
    prefix: Bytes,
    offset: usize,
}

impl<S> PrefixedIo<S> {
    fn new(inner: S, prefix: Vec<u8>) -> Self {
        Self {
            inner,
            prefix: Bytes::from(prefix),
            offset: 0,
        }
    }
}

impl<S: AsyncRead + Unpin> AsyncRead for PrefixedIo<S> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        if self.offset < self.prefix.len() && buf.remaining() > 0 {
            let remaining = &self.prefix[self.offset..];
            let n = remaining.len().min(buf.remaining());
            buf.put_slice(&remaining[..n]);
            self.offset += n;
            return Poll::Ready(Ok(()));
        }

        Pin::new(&mut self.inner).poll_read(cx, buf)
    }
}

impl<S: AsyncWrite + Unpin> AsyncWrite for PrefixedIo<S> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        data: &[u8],
    ) -> Poll<std::io::Result<usize>> {
        Pin::new(&mut self.inner).poll_write(cx, data)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.inner).poll_shutdown(cx)
    }
}

#[cfg(test)]
mod test {
    use super::{filter_h2_header, validate_h2_request_parts};
    use http::{HeaderName, HeaderValue, Request};

    #[test]
    fn test_filter_h2_forbidden_headers() {
        let name = HeaderName::from_static("connection");
        let value = HeaderValue::from_static("keep-alive");
        assert!(filter_h2_header(&name, &value).is_none());

        let name = HeaderName::from_static("te");
        let value = HeaderValue::from_static("trailers");
        assert!(filter_h2_header(&name, &value).is_some());

        let name = HeaderName::from_static("te");
        let value = HeaderValue::from_static("gzip");
        assert!(filter_h2_header(&name, &value).is_none());
    }

    #[test]
    fn test_validate_h2_request_parts_ok() {
        let req = Request::builder()
            .method("GET")
            .uri("https://example.com/")
            .body(())
            .expect("request build must succeed");
        let (parts, _) = req.into_parts();
        assert!(validate_h2_request_parts(&parts).is_ok());
    }

    #[test]
    fn test_prefixed_io_reads_prefix_before_inner() {
        use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};

        let rt = tokio::runtime::Runtime::new().expect("tokio runtime");
        rt.block_on(async {
            let (mut tx, rx) = duplex(64);
            tx.write_all(b"inner").await.expect("write inner");
            drop(tx);

            let mut io = super::PrefixedIo::new(rx, b"prefix-".to_vec());
            let mut out = Vec::new();
            io.read_to_end(&mut out).await.expect("read all");
            assert_eq!(out, b"prefix-inner");
        });
    }

    #[test]
    fn test_validate_h2_request_parts_missing_authority_and_host() {
        let req = Request::builder()
            .method("GET")
            .uri("/")
            .body(())
            .expect("request build must succeed");
        let (parts, _) = req.into_parts();
        assert!(validate_h2_request_parts(&parts).is_err());
    }

    #[test]
    fn test_validate_h2_request_parts_host_fallback() {
        let req = Request::builder()
            .method("GET")
            .uri("/")
            .header(http::header::HOST, "example.com")
            .body(())
            .expect("request build must succeed");
        let (parts, _) = req.into_parts();
        assert!(validate_h2_request_parts(&parts).is_ok());
    }
}
