use std::collections::HashMap;
use std::fs::File;
use std::io::BufReader;
use std::net::SocketAddr;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result, bail};
use bytes::{BufMut, Bytes, BytesMut};
use kafka_protocol::error::ParseResponseErrorCode;
use kafka_protocol::messages::{
    ApiVersionsRequest, RequestHeader, ResponseHeader, SaslAuthenticateRequest,
    SaslHandshakeRequest,
};
use kafka_protocol::protocol::{
    Decodable, HeaderVersion, Message, Request, StrBytes, VersionRange,
    encode_request_header_into_buffer,
};
use rustls::pki_types::{CertificateDer, PrivateKeyDer, ServerName};
use rustls::{ClientConfig as RustlsClientConfig, RootCertStore};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;
use tokio_rustls::client::TlsStream;
use tracing::{Instrument, debug, trace, trace_span};

use super::scram::ScramClient;
use super::select_api_version;
use crate::config::{SaslConfig, SaslMechanism, SecurityProtocol, TlsConfig};
use crate::constants::{API_VERSIONS_FALLBACK_VERSION, API_VERSIONS_PROBE_VERSION};

pub async fn connect_to_any_bootstrap(
    servers: &[String],
    client_id: &str,
    timeout: Duration,
    security_protocol: SecurityProtocol,
    tls: &TlsConfig,
    sasl: &SaslConfig,
) -> Result<BrokerConnection> {
    if servers.is_empty() {
        bail!("no bootstrap servers configured");
    }

    let mut last_error: Option<anyhow::Error> = None;
    for server in servers {
        match BrokerConnection::connect_with_transport(
            server,
            client_id,
            timeout,
            security_protocol,
            tls,
            sasl,
        )
        .await
        {
            Ok(conn) => return Ok(conn),
            Err(e) => {
                debug!(server = %server, error = %e, "bootstrap connection failed, trying next server");
                last_error = Some(e);
            }
        }
    }
    Err(last_error.unwrap())
}

pub struct BrokerConnection {
    stream: BrokerStream,
    next_correlation_id: i32,
    api_versions: HashMap<i16, VersionRange>,
    finalized_features: HashMap<String, i16>,
}

enum BrokerStream {
    Plain(TcpStream),
    Tls(Box<TlsStream<TcpStream>>),
}

impl BrokerStream {
    async fn write_all(&mut self, frame: &[u8]) -> Result<()> {
        match self {
            Self::Plain(stream) => stream.write_all(frame).await?,
            Self::Tls(stream) => stream.write_all(frame).await?,
        }
        Ok(())
    }

    async fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        match self {
            Self::Plain(stream) => {
                stream.read_exact(buf).await?;
            }
            Self::Tls(stream) => {
                stream.read_exact(buf).await?;
            }
        };
        Ok(())
    }
}

impl BrokerConnection {
    pub async fn connect_with_transport(
        address: &str,
        client_id: &str,
        timeout: Duration,
        security_protocol: SecurityProtocol,
        tls: &TlsConfig,
        sasl: &SaslConfig,
    ) -> Result<Self> {
        async {
            debug!(?security_protocol, "connecting to broker");
            let stream = connect_stream(address, timeout, security_protocol, tls).await?;
            let mut connection = Self {
                stream,
                next_correlation_id: 1,
                api_versions: HashMap::new(),
                finalized_features: HashMap::new(),
            };
            if security_protocol.uses_sasl() {
                connection.authenticate_sasl(client_id, sasl).await?;
            }
            connection.negotiate_versions(client_id).await?;
            debug!(
                api_keys = connection.api_versions.len(),
                finalized_features = connection.finalized_features.len(),
                ?security_protocol,
                "connected to broker"
            );
            Ok(connection)
        }
        .instrument(tracing::debug_span!(
            "broker_connect",
            %address,
            %client_id,
            timeout_ms = timeout.as_millis()
        ))
        .await
    }

    async fn authenticate_sasl(&mut self, client_id: &str, sasl: &SaslConfig) -> Result<()> {
        let response = self
            .send_request::<ApiVersionsRequest>(
                client_id,
                API_VERSIONS_FALLBACK_VERSION,
                &ApiVersionsRequest::default(),
            )
            .await
            .context("SASL ApiVersions probe failed")?;
        if let Some(error) = response.error_code.err() {
            bail!("SASL ApiVersions probe failed: {error}");
        }

        let api_versions = response
            .api_keys
            .into_iter()
            .map(|api| {
                (
                    api.api_key,
                    VersionRange {
                        min: api.min_version,
                        max: api.max_version,
                    },
                )
            })
            .collect::<HashMap<_, _>>();

        let handshake_version = api_versions
            .get(&SaslHandshakeRequest::KEY)
            .copied()
            .map(|range| {
                select_api_version(
                    SaslHandshakeRequest::KEY,
                    range,
                    SaslHandshakeRequest::VERSIONS,
                    SaslHandshakeRequest::VERSIONS.max,
                )
            })
            .transpose()?
            .unwrap_or(0);
        let authenticate_version = api_versions
            .get(&SaslAuthenticateRequest::KEY)
            .copied()
            .map(|range| {
                select_api_version(
                    SaslAuthenticateRequest::KEY,
                    range,
                    SaslAuthenticateRequest::VERSIONS,
                    SaslAuthenticateRequest::VERSIONS.max,
                )
            })
            .transpose()?;

        let mechanism = sasl.mechanism.as_str();
        let handshake =
            SaslHandshakeRequest::default().with_mechanism(StrBytes::from_static_str(mechanism));
        let response = self
            .send_request::<SaslHandshakeRequest>(client_id, handshake_version, &handshake)
            .await
            .context("SASL handshake request failed")?;
        if let Some(error) = response.error_code.err() {
            let enabled = response
                .mechanisms
                .iter()
                .map(ToString::to_string)
                .collect::<Vec<_>>()
                .join(", ");
            bail!(
                "SASL handshake failed for mechanism {mechanism}: {error}; enabled mechanisms: [{enabled}]"
            );
        }

        match sasl.mechanism {
            SaslMechanism::Plain => {
                let token = build_plain_sasl_token(sasl)?;
                if let Some(version) = authenticate_version {
                    self.send_sasl_authenticate(client_id, version, mechanism, token)
                        .await?;
                } else {
                    write_raw_sasl_token(&mut self.stream, &token).await?;
                }
            }
            SaslMechanism::ScramSha256 | SaslMechanism::ScramSha512 => {
                self.authenticate_scram(client_id, sasl, authenticate_version)
                    .await?;
            }
        }

        debug!(mechanism, "completed SASL authentication");
        Ok(())
    }

    async fn authenticate_scram(
        &mut self,
        client_id: &str,
        sasl: &SaslConfig,
        authenticate_version: Option<i16>,
    ) -> Result<()> {
        let username = sasl
            .username
            .as_ref()
            .context("SASL/SCRAM requires a username")?
            .clone();
        let password = sasl
            .password
            .as_ref()
            .context("SASL/SCRAM requires a password")?
            .clone();
        let mechanism = sasl.mechanism.as_str();
        let mut scram = ScramClient::new(sasl.mechanism, username, password)?;
        let client_first = scram.client_first_message();

        let server_first = if let Some(version) = authenticate_version {
            self.send_sasl_authenticate(client_id, version, mechanism, client_first)
                .await?
        } else {
            write_raw_sasl_token(&mut self.stream, &client_first).await?;
            read_frame(&mut self.stream).await?
        };
        let client_final = scram.handle_server_first_message(&server_first)?;

        let server_final = if let Some(version) = authenticate_version {
            self.send_sasl_authenticate(client_id, version, mechanism, client_final)
                .await?
        } else {
            write_raw_sasl_token(&mut self.stream, &client_final).await?;
            read_frame(&mut self.stream).await?
        };
        scram.handle_server_final_message(&server_final)?;
        Ok(())
    }

    async fn send_sasl_authenticate(
        &mut self,
        client_id: &str,
        version: i16,
        mechanism: &str,
        token: Vec<u8>,
    ) -> Result<Vec<u8>> {
        let request = SaslAuthenticateRequest::default().with_auth_bytes(Bytes::from(token));
        let response = self
            .send_request::<SaslAuthenticateRequest>(client_id, version, &request)
            .await
            .context("SASL authenticate request failed")?;
        if let Some(error) = response.error_code.err() {
            let message = response
                .error_message
                .as_ref()
                .map(ToString::to_string)
                .filter(|message| !message.is_empty())
                .unwrap_or_else(|| error.to_string());
            bail!("SASL authentication failed for mechanism {mechanism}: {message}");
        }
        Ok(response.auth_bytes.to_vec())
    }

    pub fn version_with_cap<Req>(&self, cap: i16) -> Result<i16>
    where
        Req: Request,
    {
        let broker_range = self
            .api_versions
            .get(&Req::KEY)
            .copied()
            .with_context(|| format!("broker did not advertise API key {}", Req::KEY))?;
        select_api_version(Req::KEY, broker_range, Req::VERSIONS, cap)
    }

    pub fn finalized_feature_level(&self, feature: &str) -> Option<i16> {
        self.finalized_features.get(feature).copied()
    }

    pub fn finalized_feature_levels(&self) -> Vec<(String, i16)> {
        let mut features = self
            .finalized_features
            .iter()
            .map(|(name, level)| (name.clone(), *level))
            .collect::<Vec<_>>();
        features.sort_by(|left, right| left.0.cmp(&right.0));
        features
    }

    async fn negotiate_versions(&mut self, client_id: &str) -> Result<()> {
        let modern_request = ApiVersionsRequest::default()
            .with_client_software_name(StrBytes::from_static_str("kafkit-client"))
            .with_client_software_version(StrBytes::from_static_str("0.2.0"));

        let response = match self
            .send_request::<ApiVersionsRequest>(
                client_id,
                API_VERSIONS_PROBE_VERSION,
                &modern_request,
            )
            .await
        {
            Ok(response) => response,
            Err(error) => {
                debug!(
                    error = %error,
                    "modern ApiVersions probe failed, retrying with fallback request"
                );
                self.send_request::<ApiVersionsRequest>(
                    client_id,
                    API_VERSIONS_FALLBACK_VERSION,
                    &ApiVersionsRequest::default(),
                )
                .await?
            }
        };

        if let Some(error) = response.error_code.err() {
            bail!("ApiVersions failed: {error}");
        }

        self.api_versions = response
            .api_keys
            .into_iter()
            .map(|api| {
                (
                    api.api_key,
                    VersionRange {
                        min: api.min_version,
                        max: api.max_version,
                    },
                )
            })
            .collect();
        self.finalized_features = response
            .finalized_features
            .into_iter()
            .map(|feature| (feature.name.to_string(), feature.max_version_level))
            .collect();

        trace!(
            api_keys = self.api_versions.len(),
            finalized_features = self.finalized_features.len(),
            "negotiated broker ApiVersions"
        );
        Ok(())
    }

    pub async fn send_request<Req>(
        &mut self,
        client_id: &str,
        version: i16,
        request: &Req,
    ) -> Result<Req::Response>
    where
        Req: Request,
    {
        let correlation_id = self.next_correlation_id;
        self.next_correlation_id += 1;
        let span = trace_span!(
            "kafka_request",
            request = std::any::type_name::<Req>(),
            api_key = Req::KEY,
            api_version = version,
            correlation_id,
            %client_id
        );

        async {
            let mut body = BytesMut::new();
            let header = RequestHeader::default()
                .with_request_api_key(Req::KEY)
                .with_request_api_version(version)
                .with_correlation_id(correlation_id)
                .with_client_id(Some(StrBytes::from_string(client_id.to_owned())));
            encode_request_header_into_buffer(&mut body, &header)?;
            request.encode(&mut body, version)?;

            trace!(request_bytes = body.len(), "encoded Kafka request");

            let mut frame = BytesMut::with_capacity(body.len() + 4);
            frame.put_i32(i32::try_from(body.len()).context("request frame is too large")?);
            frame.extend_from_slice(&body);

            self.stream.write_all(&frame).await?;
            trace!(frame_bytes = frame.len(), "wrote Kafka request frame");

            let response_frame = read_frame(&mut self.stream).await?;
            trace!(
                response_bytes = response_frame.len(),
                "received Kafka response frame"
            );
            let mut response_body = Bytes::from(response_frame);
            let header_version = Req::Response::header_version(version);
            let response_header = ResponseHeader::decode(&mut response_body, header_version)?;
            if response_header.correlation_id != correlation_id {
                bail!(
                    "response correlation mismatch: expected {}, got {}",
                    correlation_id,
                    response_header.correlation_id
                );
            }

            let response = Req::Response::decode(&mut response_body, version)?;
            debug!("completed Kafka request");
            Ok(response)
        }
        .instrument(span)
        .await
    }
}

async fn connect_stream(
    address: &str,
    timeout: Duration,
    security_protocol: SecurityProtocol,
    tls: &TlsConfig,
) -> Result<BrokerStream> {
    let tcp_stream = tokio::time::timeout(timeout, TcpStream::connect(address))
        .await
        .with_context(|| format!("timed out connecting to {address}"))?
        .with_context(|| format!("failed to connect to {address}"))?;
    tcp_stream
        .set_nodelay(true)
        .with_context(|| format!("failed to enable TCP_NODELAY for {address}"))?;

    if security_protocol.uses_tls() {
        let tls_config = build_tls_client_config(tls)?;
        let connector = TlsConnector::from(tls_config);
        let server_name = server_name_for_tls(address, tls)?;
        let stream = tokio::time::timeout(timeout, connector.connect(server_name, tcp_stream))
            .await
            .with_context(|| format!("timed out negotiating TLS with {address}"))?
            .with_context(|| format!("failed TLS handshake with {address}"))?;
        Ok(BrokerStream::Tls(Box::new(stream)))
    } else {
        Ok(BrokerStream::Plain(tcp_stream))
    }
}

fn build_plain_sasl_token(sasl: &SaslConfig) -> Result<Vec<u8>> {
    let username = sasl
        .username
        .as_deref()
        .context("SASL/PLAIN requires a username")?;
    let password = sasl
        .password
        .as_deref()
        .context("SASL/PLAIN requires a password")?;
    let authorization_id = sasl.authorization_id.as_deref().unwrap_or_default();

    let mut token =
        Vec::with_capacity(authorization_id.len() + username.len() + password.len() + 2);
    token.extend_from_slice(authorization_id.as_bytes());
    token.push(0);
    token.extend_from_slice(username.as_bytes());
    token.push(0);
    token.extend_from_slice(password.as_bytes());
    Ok(token)
}

async fn write_raw_sasl_token(stream: &mut BrokerStream, token: &[u8]) -> Result<()> {
    let mut frame = BytesMut::with_capacity(token.len() + 4);
    frame.put_i32(i32::try_from(token.len()).context("SASL token frame is too large")?);
    frame.extend_from_slice(token);
    stream.write_all(&frame).await
}

fn build_tls_client_config(tls: &TlsConfig) -> Result<Arc<RustlsClientConfig>> {
    let mut root_store = RootCertStore::empty();
    root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());

    if let Some(ca_cert_path) = tls.ca_cert_path.as_deref() {
        for cert in load_certificates(ca_cert_path)? {
            root_store.add(cert)?;
        }
    }

    let builder = RustlsClientConfig::builder().with_root_certificates(root_store);
    let config = match (
        tls.client_cert_path.as_deref(),
        tls.client_key_path.as_deref(),
    ) {
        (Some(client_cert_path), Some(client_key_path)) => builder.with_client_auth_cert(
            load_certificates(client_cert_path)?,
            load_private_key(client_key_path)?,
        )?,
        (None, None) => builder.with_no_client_auth(),
        _ => bail!("TLS client auth requires both client_cert_path and client_key_path"),
    };

    Ok(Arc::new(config))
}

fn load_certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>> {
    let file = File::open(path)
        .with_context(|| format!("failed to open TLS certificate file '{}'", path.display()))?;
    let mut reader = BufReader::new(file);
    let certs = rustls_pemfile::certs(&mut reader)
        .collect::<std::result::Result<Vec<_>, _>>()
        .with_context(|| format!("failed to parse TLS certificate PEM '{}'", path.display()))?;
    if certs.is_empty() {
        bail!(
            "TLS certificate file '{}' did not contain any PEM certificates",
            path.display()
        );
    }
    Ok(certs)
}

fn load_private_key(path: &Path) -> Result<PrivateKeyDer<'static>> {
    let file = File::open(path)
        .with_context(|| format!("failed to open TLS private key file '{}'", path.display()))?;
    let mut reader = BufReader::new(file);
    rustls_pemfile::private_key(&mut reader)
        .with_context(|| format!("failed to parse TLS private key PEM '{}'", path.display()))?
        .with_context(|| {
            format!(
                "TLS private key file '{}' did not contain a PEM key",
                path.display()
            )
        })
}

fn server_name_for_tls(address: &str, tls: &TlsConfig) -> Result<ServerName<'static>> {
    if let Some(server_name) = tls.server_name.as_ref() {
        return ServerName::try_from(server_name.clone())
            .with_context(|| format!("invalid TLS server name '{}'", server_name));
    }

    if let Ok(socket_addr) = address.parse::<SocketAddr>() {
        return Ok(ServerName::IpAddress(socket_addr.ip().into()));
    }

    let host = if let Some(stripped) = address.strip_prefix('[') {
        stripped
            .split(']')
            .next()
            .context("invalid bracketed broker address")?
            .to_owned()
    } else {
        address
            .rsplit_once(':')
            .map(|(host, _)| host.to_owned())
            .unwrap_or_else(|| address.to_owned())
    };

    ServerName::try_from(host.clone()).with_context(|| {
        format!("could not derive a valid TLS server name from broker address '{address}'")
    })
}

async fn read_frame(stream: &mut BrokerStream) -> Result<Vec<u8>> {
    let mut header = [0_u8; 4];
    stream.read_exact(&mut header).await?;
    let frame_len = i32::from_be_bytes(header);
    if frame_len < 0 {
        bail!("broker returned a negative frame length: {frame_len}");
    }

    let mut payload = vec![0_u8; usize::try_from(frame_len)?];
    stream.read_exact(&mut payload).await?;
    Ok(payload)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    #[test]
    fn tls_server_name_defaults_to_host() {
        let server_name =
            server_name_for_tls("broker.example.com:9093", &TlsConfig::default()).unwrap();
        assert_eq!(server_name.to_str(), "broker.example.com");
    }

    #[test]
    fn tls_server_name_respects_explicit_override() {
        let tls = TlsConfig::new().with_server_name("cluster.internal");
        let server_name = server_name_for_tls("127.0.0.1:9093", &tls).unwrap();
        assert_eq!(server_name.to_str(), "cluster.internal");
    }

    #[test]
    fn tls_server_name_handles_ip_and_bracketed_ipv6() {
        let server_name = server_name_for_tls("127.0.0.1:9093", &TlsConfig::default()).unwrap();
        assert_eq!(server_name.to_str(), "127.0.0.1");

        let server_name = server_name_for_tls("[::1]:9093", &TlsConfig::default()).unwrap();
        assert_eq!(server_name.to_str(), "::1");
    }

    #[test]
    fn tls_server_name_rejects_invalid_override_and_empty_address() {
        let tls = TlsConfig::new().with_server_name("not a dns name");
        assert!(server_name_for_tls("127.0.0.1:9093", &tls).is_err());
        assert!(server_name_for_tls("", &TlsConfig::default()).is_err());
    }

    #[test]
    fn plain_sasl_token_requires_credentials_and_uses_authorization_id() {
        assert!(build_plain_sasl_token(&SaslConfig::default()).is_err());
        assert!(
            build_plain_sasl_token(&SaslConfig::plain("user", "pw").with_authorization_id("authz"))
                .unwrap()
                == b"authz\0user\0pw"
        );
    }

    #[test]
    fn tls_file_loaders_reject_missing_empty_and_invalid_pem_files() {
        let dir =
            std::env::temp_dir().join(format!("kafkit-client-tls-test-{}", std::process::id()));
        fs::create_dir_all(&dir).unwrap();
        let cert_path = dir.join("cert.pem");
        let key_path = dir.join("key.pem");
        fs::write(&cert_path, b"not a certificate").unwrap();
        fs::write(&key_path, b"not a key").unwrap();

        assert!(load_certificates(&cert_path).is_err());
        assert!(load_private_key(&key_path).is_err());
        assert!(load_certificates(&dir.join("missing.pem")).is_err());

        let _ = fs::remove_dir_all(dir);
    }
}
