use std::io::Read;
use std::io::Write;
use std::net::Shutdown;
use std::net::TcpStream;
#[cfg(unix)]
use std::os::unix::net::UnixStream;
use std::time::Duration;

use qroissant_core::Attribute;
use qroissant_core::Compression;
use qroissant_core::Encoding;
use qroissant_core::Frame;
use qroissant_core::HEADER_LEN;
use qroissant_core::MessageHeader;
use qroissant_core::MessageType;
use qroissant_core::StreamingDecompressor;
use qroissant_core::Value;
use qroissant_core::Vector;
use qroissant_core::VectorData;
use qroissant_core::encode_message;
use qroissant_core::read_frame;
use qroissant_core::read_message_length;

use crate::TransportError;
use crate::TransportResult;

pub const CLIENT_CAPABILITY: u8 = 3;

pub enum SyncTransport {
    Tcp(TcpStream),
    #[cfg(unix)]
    Unix(UnixStream),
}

impl SyncTransport {
    pub fn shutdown(&mut self) -> std::io::Result<()> {
        match self {
            Self::Tcp(stream) => stream.shutdown(Shutdown::Both),
            #[cfg(unix)]
            Self::Unix(stream) => stream.shutdown(Shutdown::Both),
        }
    }

    pub fn take_error(&self) -> std::io::Result<Option<std::io::Error>> {
        match self {
            Self::Tcp(stream) => stream.take_error(),
            #[cfg(unix)]
            Self::Unix(stream) => stream.take_error(),
        }
    }

    pub fn set_timeouts(&self, timeout_ms: Option<u64>) -> std::io::Result<()> {
        let timeout = timeout_ms.map(Duration::from_millis);
        match self {
            Self::Tcp(stream) => {
                stream.set_read_timeout(timeout)?;
                stream.set_write_timeout(timeout)?;
                stream.set_nodelay(true)
            }
            #[cfg(unix)]
            Self::Unix(stream) => {
                stream.set_read_timeout(timeout)?;
                stream.set_write_timeout(timeout)
            }
        }
    }
}

impl Read for SyncTransport {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        match self {
            Self::Tcp(stream) => stream.read(buf),
            #[cfg(unix)]
            Self::Unix(stream) => stream.read(buf),
        }
    }
}

impl Write for SyncTransport {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        match self {
            Self::Tcp(stream) => stream.write(buf),
            #[cfg(unix)]
            Self::Unix(stream) => stream.write(buf),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Self::Tcp(stream) => stream.flush(),
            #[cfg(unix)]
            Self::Unix(stream) => stream.flush(),
        }
    }
}

pub struct SyncPooledTransport {
    transport: SyncTransport,
    broken: bool,
}

impl SyncPooledTransport {
    pub fn new(transport: SyncTransport) -> Self {
        Self {
            transport,
            broken: false,
        }
    }

    pub fn mark_broken(&mut self) {
        self.broken = true;
    }

    pub fn is_broken(&self) -> bool {
        self.broken || self.transport.take_error().ok().flatten().is_some()
    }

    pub fn shutdown(&mut self) -> std::io::Result<()> {
        self.transport.shutdown()
    }
}

impl Read for SyncPooledTransport {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.transport.read(buf)
    }
}

impl Write for SyncPooledTransport {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.transport.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.transport.flush()
    }
}

pub fn credentials_bytes(username: Option<&str>, password: Option<&str>) -> Vec<u8> {
    let username = username.unwrap_or_default();
    let password = password.unwrap_or_default();
    let mut bytes = format!("{username}:{password}").into_bytes();
    bytes.push(CLIENT_CAPABILITY);
    bytes.push(0);
    bytes
}

pub fn perform_handshake<S: Read + Write>(
    stream: &mut S,
    username: Option<&str>,
    password: Option<&str>,
) -> TransportResult<u8> {
    stream.write_all(&credentials_bytes(username, password))?;
    stream.flush()?;

    let mut capability = [0_u8; 1];
    stream.read_exact(&mut capability)?;
    Ok(capability[0])
}

pub fn encode_sync_query(message: &str) -> TransportResult<Vec<u8>> {
    let _ = i32::try_from(message.len())
        .map_err(|_| TransportError::InvalidQueryLength(message.len()))?;
    let value = Value::Vector(Vector::new(
        Attribute::None,
        VectorData::Char(bytes::Bytes::copy_from_slice(message.as_bytes())),
    ));

    encode_message(
        &value,
        Encoding::LittleEndian,
        MessageType::Synchronous,
        Compression::Uncompressed,
    )
    .map_err(|error| TransportError::Protocol(error.to_string()))
}

pub fn extract_q_error(frame_bytes: &[u8]) -> TransportResult<Option<String>> {
    let frame =
        Frame::parse(frame_bytes).map_err(|error| TransportError::Protocol(error.to_string()))?;
    let body = frame.body();
    if body.first().copied() != Some(128) {
        return Ok(None);
    }

    let message = match body[1..].iter().position(|byte| *byte == 0) {
        Some(end) => &body[1..1 + end],
        None => &body[1..],
    };
    Ok(Some(String::from_utf8_lossy(message).into_owned()))
}

pub fn parse_message_header(header_bytes: [u8; HEADER_LEN]) -> TransportResult<MessageHeader> {
    MessageHeader::from_bytes(header_bytes)
        .map_err(|error| TransportError::Protocol(error.to_string()))
}

pub fn validate_response_header(header: MessageHeader) -> TransportResult<()> {
    if header.message_type() != MessageType::Response {
        return Err(TransportError::Protocol(format!(
            "expected a q response frame, received {:?}",
            header.message_type()
        )));
    }

    Ok(())
}

pub fn validate_response_header_bytes(
    header_bytes: [u8; HEADER_LEN],
) -> TransportResult<MessageHeader> {
    let header = parse_message_header(header_bytes)?;
    validate_response_header(header)?;
    Ok(header)
}

pub fn validate_response_frame(frame_bytes: &[u8]) -> TransportResult<MessageHeader> {
    let frame =
        Frame::parse(frame_bytes).map_err(|error| TransportError::Protocol(error.to_string()))?;
    let header = frame.header();
    validate_response_header(header)?;
    Ok(header)
}

pub fn connect_tcp_transport(
    host: &str,
    port: u16,
    username: Option<&str>,
    password: Option<&str>,
    timeout_ms: Option<u64>,
) -> TransportResult<SyncTransport> {
    let mut stream = SyncTransport::Tcp(TcpStream::connect((host, port))?);
    stream.set_timeouts(timeout_ms)?;
    perform_handshake(&mut stream, username, password)?;
    Ok(stream)
}

#[cfg(unix)]
pub fn connect_unix_transport(
    path: &str,
    username: Option<&str>,
    password: Option<&str>,
    timeout_ms: Option<u64>,
) -> TransportResult<SyncTransport> {
    let mut stream = SyncTransport::Unix(UnixStream::connect(path)?);
    stream.set_timeouts(timeout_ms)?;
    perform_handshake(&mut stream, username, password)?;
    Ok(stream)
}

pub fn request_frame_over<S: Read + Write>(
    stream: &mut S,
    payload: &[u8],
) -> TransportResult<Vec<u8>> {
    stream.write_all(payload)?;
    stream.flush()?;
    read_frame(stream).map_err(|error| TransportError::Protocol(error.to_string()))
}

/// Sends a payload and reads the response frame, using streaming decompression
/// when the response is compressed.
///
/// For compressed frames, the body is read in chunks and fed to a
/// [`StreamingDecompressor`] incrementally, overlapping network I/O with
/// decompression work.  The returned frame is reconstructed as an
/// *uncompressed* frame so callers can decode it normally.
///
/// For uncompressed frames, this behaves identically to [`request_frame_over`].
pub fn request_frame_streaming_over<S: Read + Write>(
    stream: &mut S,
    payload: &[u8],
) -> TransportResult<Vec<u8>> {
    stream.write_all(payload)?;
    stream.flush()?;

    // Read the 8-byte header.
    let mut header_bytes = [0_u8; HEADER_LEN];
    stream.read_exact(&mut header_bytes)?;

    let header = parse_message_header(header_bytes)?;
    let body_len = header.body_len();

    if header.compression() == Compression::Uncompressed {
        // Fast path: read entire uncompressed body.
        let mut frame = vec![0_u8; header.size()];
        frame[..HEADER_LEN].copy_from_slice(&header_bytes);
        stream.read_exact(&mut frame[HEADER_LEN..])?;
        return Ok(frame);
    }

    // Compressed frame: read the 4-byte size prefix first.
    if body_len < 4 {
        return Err(TransportError::Protocol(
            "compressed body must be at least 4 bytes for size prefix".to_string(),
        ));
    }

    let mut size_prefix = [0_u8; 4];
    stream.read_exact(&mut size_prefix)?;

    let mut decompressor = StreamingDecompressor::new(size_prefix, header.encoding())
        .map_err(|error| TransportError::Protocol(error.to_string()))?;

    // Read the remaining compressed body in chunks.
    let remaining = body_len - 4;
    let mut total_read = 0_usize;
    let mut chunk = [0_u8; 8192];

    while total_read < remaining {
        let want = (remaining - total_read).min(chunk.len());
        stream.read_exact(&mut chunk[..want])?;
        decompressor
            .feed(&chunk[..want])
            .map_err(|error| TransportError::Protocol(error.to_string()))?;
        total_read += want;
    }

    if !decompressor.is_complete() {
        return Err(TransportError::Protocol(
            "streaming decompression did not complete after reading entire body".to_string(),
        ));
    }

    let decompressed = decompressor
        .finish()
        .map_err(|error| TransportError::Protocol(error.to_string()))?;

    // Reconstruct as an uncompressed frame: header + decompressed body.
    let new_size = HEADER_LEN + decompressed.len();
    let new_header = MessageHeader::new(
        header.encoding(),
        header.message_type(),
        Compression::Uncompressed,
        new_size,
    )
    .map_err(|error| TransportError::Protocol(error.to_string()))?;

    let mut frame = Vec::with_capacity(new_size);
    frame.extend_from_slice(
        &new_header
            .to_bytes()
            .map_err(|error| TransportError::Protocol(error.to_string()))?,
    );
    frame.extend_from_slice(&decompressed);
    Ok(frame)
}

pub fn begin_streaming_frame_over<S: Read + Write>(
    stream: &mut S,
    payload: &[u8],
) -> TransportResult<([u8; HEADER_LEN], usize)> {
    stream.write_all(payload)?;
    stream.flush()?;

    let mut header = [0_u8; HEADER_LEN];
    stream.read_exact(&mut header)?;
    let message_length = read_message_length(&header)
        .map_err(|error| TransportError::Protocol(error.to_string()))?;
    Ok((header, message_length - HEADER_LEN))
}

pub struct SyncConnection {
    transport: Option<SyncTransport>,
}

impl SyncConnection {
    pub fn connect_tcp(
        host: &str,
        port: u16,
        username: Option<&str>,
        password: Option<&str>,
        timeout_ms: Option<u64>,
    ) -> TransportResult<Self> {
        Ok(Self {
            transport: Some(connect_tcp_transport(
                host, port, username, password, timeout_ms,
            )?),
        })
    }

    #[cfg(unix)]
    pub fn connect_unix(
        path: &str,
        username: Option<&str>,
        password: Option<&str>,
        timeout_ms: Option<u64>,
    ) -> TransportResult<Self> {
        Ok(Self {
            transport: Some(connect_unix_transport(
                path, username, password, timeout_ms,
            )?),
        })
    }

    pub fn query_frame(&mut self, message: &str) -> TransportResult<Vec<u8>> {
        let payload = encode_sync_query(message)?;
        let transport = self.transport.as_mut().ok_or(TransportError::Closed)?;
        let frame = request_frame_over(transport, &payload)?;
        validate_response_frame(&frame)?;
        Ok(frame)
    }

    pub fn is_closed(&self) -> bool {
        self.transport.is_none()
    }

    pub fn close(&mut self) -> TransportResult<()> {
        let Some(mut transport) = self.transport.take() else {
            return Ok(());
        };
        transport.shutdown()?;
        Ok(())
    }
}

impl Drop for SyncConnection {
    fn drop(&mut self) {
        let _ = self.close();
    }
}
