//! WSGI environment building and response parsing

use std::str::FromStr;

use super::{ResponseState, WsgiInput};
use crate::errors::{FalcornError, Result};
use crate::http::{Body, HttpRequest, HttpResponse};
use crate::python::wsgi::WsgiContext;

use http::{HeaderMap, StatusCode};
use pyo3::prelude::*;
use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList, PyTuple};
use std::sync::Arc;

/// Builds WSGI environ dictionary from HTTP request
pub struct WsgiEnvironBuilder;

impl WsgiEnvironBuilder {
    /// Create WSGI environ dict from HTTP request
    pub fn from_request(py: Python, request: &HttpRequest) -> Result<Py<PyDict>> {
        let ctx = WsgiContext::global();
        let environ = ctx.take_environ(py)?;
        let environ_ref = environ.bind(py);

        // Request method
        environ_ref
            .set_item(
                ctx.get_or_insert_str("REQUEST_METHOD", py),
                request.method.as_str(),
            )
            .map_err(|_| FalcornError::config("Failed to set REQUEST_METHOD"))?;

        // Path info
        let path_str = std::str::from_utf8(&request.path)
            .map_err(|_| FalcornError::config("Invalid PATH_INFO encoding"))?;
        environ_ref
            .set_item("PATH_INFO", path_str)
            .map_err(|_| FalcornError::config("Failed to set PATH_INFO"))?;

        // Query string
        if let Some(qs) = &request.query_string {
            let qs_str = std::str::from_utf8(qs)
                .map_err(|_| FalcornError::config("Invalid QUERY_STRING encoding"))?;
            environ_ref
                .set_item(ctx.get_or_insert_str("QUERY_STRING", py), qs_str)
                .map_err(|_| FalcornError::config("Failed to set QUERY_STRING"))?;
        } else {
            environ_ref
                .set_item(ctx.get_or_insert_str("QUERY_STRING", py), "")
                .map_err(|_| FalcornError::config("Failed to set REQUEST_METHOD"))?;
        }

        // Content type and length
        if let Some(ct) = request.get_header("content-type") {
            let ct_str = std::str::from_utf8(ct)
                .map_err(|_| FalcornError::config("Invalid CONTENT_TYPE encoding"))?;
            environ_ref
                .set_item(ctx.get_or_insert_str("CONTENT_TYPE", py), ct_str)
                .map_err(|_| FalcornError::config("Failed to set CONTENT_TYPE"))?;
        }

        if let Some(cl) = request.content_length() {
            environ_ref
                .set_item(ctx.get_or_insert_str("CONTENT_LENGTH", py), cl.to_string())
                .map_err(|_| FalcornError::config("Failed to set CONTENT_LENGTH"))?;
        }

        // Server info
        let addr = request.peer_addr.to_string();
        let parts: Vec<&str> = addr.split(':').collect();
        if parts.len() >= 1 {
            environ_ref
                .set_item(ctx.get_or_insert_str("REMOTE_ADDR", py), parts[0])
                .map_err(|_| FalcornError::config("Failed to set REMOTE_ADDR"))?;
        }

        // HTTP version
        environ_ref
            .set_item(
                ctx.get_or_insert_str("SERVER_PROTOCOL", py),
                &request.version_str(),
            )
            .map_err(|_| FalcornError::config("Failed to set SERVER_PROTOCOL"))?;

        // Server host and port
        if let Some(host) = request.get_header("host") {
            if let Some(idx) = host.iter().position(|b| *b == b':') {
                if let Ok(host_str) = std::str::from_utf8(&host[..idx]) {
                    environ_ref
                        .set_item(ctx.get_or_insert_str("SERVER_NAME", py), host_str)
                        .map_err(|_| FalcornError::config("Failed to set SERVER_NAME"))?;
                }
                if let Ok(port_str) = std::str::from_utf8(&host[idx + 1..]) {
                    environ_ref
                        .set_item(ctx.get_or_insert_str("SERVER_PORT", py), port_str)
                        .map_err(|_| FalcornError::config("Failed to set SERVER_PORT"))?;
                }
            } else if let Ok(host_str) = std::str::from_utf8(host) {
                environ_ref
                    .set_item(ctx.get_or_insert_str("SERVER_NAME", py), host_str)
                    .map_err(|_| FalcornError::config("Failed to set SERVER_NAME"))?;
            }
        }

        // Script name (already set in base, but allow override)
        environ_ref.set_item(ctx.get_or_insert_str("SCRIPT_NAME", py), "")?;

        // Add all headers with HTTP_ prefix
        for (name, value) in &request.headers {
            if let Ok(value_str) = std::str::from_utf8(value) {
                let key = ctx.get_or_insert_http_header_key(name, py);
                environ_ref
                    .set_item(key, value_str)
                    .map_err(|_| FalcornError::config("Failed to set header"))?;
            }
        }

        // Request body
        if let Some(body) = &request.body {
            let input = WsgiInput::new(body.clone());
            let py_input = Py::new(py, input)?;
            environ_ref
                .set_item(ctx.get_or_insert_str("wsgi.input", py), py_input)
                .map_err(|_| FalcornError::config("Failed to set wsgi.input"))?;
        } else {
            let empty_input = Py::new(py, WsgiInput::empty_input())?;
            environ_ref
                .set_item(ctx.get_or_insert_str("wsgi.input", py), empty_input)
                .map_err(|_| FalcornError::config("Failed to set wsgi.input"))?;
        }

        // wsgi.* keys are already set in base environ

        Ok(environ)
    }
}

/// Response adapter for WSGI applications
pub struct WsgiResponseAdapter {
    pub status: StatusCode,
    pub headers: HeaderMap,
    pub body_chunks: Vec<bytes::Bytes>,
}

impl WsgiResponseAdapter {
    /// Create from WSGI application output
    pub fn from_wsgi_output<'py>(
        _py: Python<'py>,
        app_output: &Bound<'py, PyAny>,
        state: Arc<ResponseState>,
    ) -> Result<Self> {
        // WSGI app returns (status, headers, body_iterable)

        // Extract status and header
        // Status and headers should have been set by start_response
        // They are stored in thread-local storage by the callback
        let (status, headers) = Self::extract_header_and_status(state)?;

        // Fast-path body extraction
        let mut body_chunks: Vec<bytes::Bytes> = Vec::new();

        // Bytes Output
        if app_output.is_instance_of::<PyBytes>() {
            let bytes = app_output
                .extract::<&[u8]>()
                .map_err(|e| FalcornError::runtime(e.to_string()))?;
            if !bytes.is_empty() {
                body_chunks.push(bytes::Bytes::copy_from_slice(bytes));
            }
        }
        // ByteArray Output
        else if app_output.is_instance_of::<PyByteArray>() {
            let bytes = app_output
                .extract::<&[u8]>()
                .map_err(|e| FalcornError::runtime(e.to_string()))?;
            if !bytes.is_empty() {
                body_chunks.push(bytes::Bytes::copy_from_slice(bytes));
            }
        }
        // List Output
        else if app_output.is_instance_of::<PyList>() {
            let list = app_output
                .cast::<PyList>()
                .map_err(|e| FalcornError::runtime(e.to_string()))?;
            for item in list.iter() {
                let bytes = item
                    .extract::<&[u8]>()
                    .map_err(|_| FalcornError::PythonTypeError {
                        reason: "Invalid body chunk type".to_string(),
                    })?;
                if !bytes.is_empty() {
                    body_chunks.push(bytes::Bytes::copy_from_slice(bytes));
                }
            }
        }
        // Tuple Output
        else if app_output.is_instance_of::<PyTuple>() {
            let tuple =
                app_output
                    .cast::<PyTuple>()
                    .map_err(|_| FalcornError::PythonTypeError {
                        reason: "Invalid body chunk type".to_string(),
                    })?;
            for item in tuple.iter() {
                let bytes = item
                    .extract::<&[u8]>()
                    .map_err(|_| FalcornError::PythonTypeError {
                        reason: "Invalid body chunk type".to_string(),
                    })?;
                if !bytes.is_empty() {
                    body_chunks.push(bytes::Bytes::copy_from_slice(bytes));
                }
            }
        } else {
            // Fallback: iterate any iterable
            let body_iter = app_output.try_iter()?;
            for chunk in body_iter {
                let chunk = chunk?;
                let bytes =
                    chunk
                        .extract::<&[u8]>()
                        .map_err(|_| FalcornError::PythonTypeError {
                            reason: "Invalid body chunk type".to_string(),
                        })?;
                if !bytes.is_empty() {
                    body_chunks.push(bytes::Bytes::copy_from_slice(bytes));
                }
            }
        }

        Ok(Self {
            status,
            headers,
            body_chunks,
        })
    }

    /// Convert to HTTP response
    pub fn to_http_response(&self) -> Result<HttpResponse> {
        use crate::http::response;

        let body = match self.body_chunks.len() {
            0 => Body::empty(),
            1 => Body::from_bytes(self.body_chunks[0].clone()),
            _ => Body::from_sync_chunks(self.body_chunks.clone()),
        };

        // Avoid ResponseBuilder for streaming bodies to prevent incorrect Content-Length.
        if matches!(
            body,
            Body::SyncReader(_) | Body::AsyncStream(_) | Body::File { .. }
        ) {
            return Ok(HttpResponse {
                status: self.status,
                headers: self.headers.clone(),
                body,
                version: http::Version::HTTP_11,
            });
        }

        let response = response()
            // set status
            .status(self.status)
            // headers
            .set_headers(self.headers.clone())
            // Body
            .body(match &body {
                Body::Bytes(bytes) => bytes.to_vec(),
                Body::Empty => Vec::new(),
                Body::SyncReader(_) | Body::AsyncStream(_) | Body::File { .. } => Vec::new(),
            });

        let mut resp = response.build();
        resp.body = body;
        Ok(resp)
    }

    /// Get response data from thread-local storage
    fn extract_header_and_status(state: Arc<ResponseState>) -> Result<(StatusCode, HeaderMap)> {
        let mut status_code = state
            .status
            .pop()
            .ok_or_else(|| FalcornError::runtime("start_response not called"))?;

        let code = status_code.split_off(3);
        drop(code);

        let mut headers_vec = Vec::new();

        while let Some(h) = state.headers.pop() {
            headers_vec.push(h);
        }

        // Convert to StatusCode
        let status = StatusCode::from_str(&status_code)
            .map_err(|_| FalcornError::http_parse("Invalid status code"))?;

        // Convert headers
        let mut headers = HeaderMap::new();
        for item in headers_vec.iter() {
            let name: String = item.0.clone();
            let value: String = item.1.clone();

            if let (Ok(header_name), Ok(header_value)) = (
                http::header::HeaderName::try_from(name.as_str()),
                http::header::HeaderValue::try_from(value.as_str()),
            ) {
                headers.insert(header_name, header_value);
            }
        }

        Ok((status, headers))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::http::HttpMethod;
    use bytes::Bytes;

    fn sample_request_with_body() -> HttpRequest {
        HttpRequest {
            raw: Bytes::new(),
            method: HttpMethod::Post,
            path: Bytes::from_static(b"/submit"),
            query_string: Some(Bytes::from_static(b"a=1")),
            version: http::Version::HTTP_11,
            headers: vec![
                (
                    Bytes::from_static(b"host"),
                    Bytes::from_static(b"example.com:8080"),
                ),
                (
                    Bytes::from_static(b"content-type"),
                    Bytes::from_static(b"text/plain"),
                ),
                (
                    Bytes::from_static(b"content-length"),
                    Bytes::from_static(b"5"),
                ),
                (
                    Bytes::from_static(b"x-request-id"),
                    Bytes::from_static(b"abc123"),
                ),
            ],
            body: Some(Arc::new(Bytes::from_static(b"hello"))),
            peer_addr: "127.0.0.1:43210".parse().expect("socket addr"),
        }
    }

    #[test]
    fn test_wsgi_environ_builder_populates_common_fields() {
        let request = sample_request_with_body();

        Python::attach(|py| {
            let environ = WsgiEnvironBuilder::from_request(py, &request).expect("environ");
            let env = environ.bind(py);

            assert_eq!(
                env.get_item("REQUEST_METHOD")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "POST"
            );
            assert_eq!(
                env.get_item("PATH_INFO")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "/submit"
            );
            assert_eq!(
                env.get_item("QUERY_STRING")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "a=1"
            );
            assert_eq!(
                env.get_item("SERVER_NAME")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "example.com"
            );
            assert_eq!(
                env.get_item("SERVER_PORT")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "8080"
            );
            assert_eq!(
                env.get_item("CONTENT_TYPE")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "text/plain"
            );
            assert_eq!(
                env.get_item("CONTENT_LENGTH")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "5"
            );
            assert_eq!(
                env.get_item("HTTP_X_REQUEST_ID")
                    .expect("get")
                    .expect("value")
                    .extract::<String>()
                    .expect("str"),
                "abc123"
            );

            let wsgi_input = env
                .get_item("wsgi.input")
                .expect("get")
                .expect("value")
                .clone();
            let body = wsgi_input
                .call_method1("read", (5,))
                .expect("read")
                .extract::<Vec<u8>>()
                .expect("extract");
            assert_eq!(body, b"hello");
        });
    }

    #[test]
    fn test_wsgi_response_adapter_from_list_output() {
        Python::attach(|py| {
            let state = Arc::new(ResponseState::new());
            state.status.push("200 OK".to_string());
            state
                .headers
                .push(("content-type".to_string(), "text/plain".to_string()));

            let output =
                PyList::new(py, [PyBytes::new(py, b"he"), PyBytes::new(py, b"llo")]).expect("list");
            let adapted = WsgiResponseAdapter::from_wsgi_output(py, &output.into_any(), state)
                .expect("adapt");

            assert_eq!(adapted.status, StatusCode::OK);
            let combined: Vec<u8> = adapted
                .body_chunks
                .iter()
                .flat_map(|b| b.as_ref().iter().copied())
                .collect();
            assert_eq!(combined, b"hello");
            assert_eq!(
                adapted
                    .headers
                    .get("content-type")
                    .and_then(|v| v.to_str().ok()),
                Some("text/plain")
            );
        });
    }

    #[test]
    fn test_wsgi_response_adapter_requires_start_response() {
        Python::attach(|py| {
            let state = Arc::new(ResponseState::new());
            let output = PyBytes::new(py, b"body");
            let err = match WsgiResponseAdapter::from_wsgi_output(py, &output.into_any(), state) {
                Ok(_) => panic!("adapter should fail when start_response was not called"),
                Err(err) => err,
            };
            assert!(matches!(err, FalcornError::RuntimeError { .. }));
        });
    }
}
