use std::collections::HashMap;
use std::fmt;
use std::sync::atomic::{AtomicBool, AtomicU16, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use rustls::client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier};
use rustls::crypto::{ring, verify_tls12_signature, verify_tls13_signature, CryptoProvider};
use rustls::pki_types::{CertificateDer, ServerName, UnixTime};
use rustls::{ClientConfig as RustlsConfig, DigitallySignedStruct, RootCertStore, SignatureScheme};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::{broadcast, mpsc, oneshot, watch};
use tokio::task::JoinHandle;
use tokio::time::timeout;
use tokio_rustls::TlsConnector;

use crate::codec::{decode_frame_packet, encode_request, CodecError};
use crate::frame::{FrameDecoder, FrameError};
use crate::packet::{command, Packet};
use crate::proxy::{connect as connect_tcp, ProxyConfig};

const READ_BUFFER_SIZE: usize = 64 * 1024;
const PUSH_CAPACITY: usize = 256;
const MINCIFRY_CA_PEM: &str = include_str!("mincifry_ca.pem");

static TRUST_MINCIFRY_CA: AtomicBool = AtomicBool::new(false);

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    Out,
    In,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WireEvent {
    pub direction: Direction,
    pub packet: Packet,
}

#[derive(Debug)]
pub enum TransportError {
    Closed,
    Timeout,
    ConnectTimeout,
    Io(std::io::Error),
    Tls(String),
    InvalidConfig(String),
    Codec(CodecError),
    Frame(FrameError),
    Server {
        message: String,
        key: Option<String>,
    },
}

impl fmt::Display for TransportError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Closed => formatter.write_str("connection is closed"),
            Self::Timeout => formatter.write_str("request timed out"),
            Self::ConnectTimeout => formatter.write_str("connection timed out"),
            Self::Io(error) => error.fmt(formatter),
            Self::Tls(message) => write!(formatter, "TLS error: {message}"),
            Self::InvalidConfig(message) => write!(formatter, "invalid configuration: {message}"),
            Self::Codec(error) => error.fmt(formatter),
            Self::Frame(error) => error.fmt(formatter),
            Self::Server { message, .. } => write!(formatter, "server error: {message}"),
        }
    }
}

impl std::error::Error for TransportError {}

impl From<std::io::Error> for TransportError {
    fn from(error: std::io::Error) -> Self {
        Self::Io(error)
    }
}

impl From<CodecError> for TransportError {
    fn from(error: CodecError) -> Self {
        Self::Codec(error)
    }
}

impl From<FrameError> for TransportError {
    fn from(error: FrameError) -> Self {
        Self::Frame(error)
    }
}

#[derive(Clone, Debug)]
pub struct ClientConfig {
    pub host: String,
    pub port: u16,
    pub insecure_tls: bool,
    pub connect_timeout: Duration,
    pub request_timeout: Duration,
    pub proxy: Option<ProxyConfig>,
}

impl ClientConfig {
    pub fn new(host: impl Into<String>, port: u16) -> Self {
        Self {
            host: host.into(),
            port,
            insecure_tls: false,
            connect_timeout: Duration::from_secs(15),
            request_timeout: Duration::from_secs(30),
            proxy: None,
        }
    }
}

type PendingResult = Result<Packet, TransportError>;
type PendingMap = Arc<Mutex<HashMap<u16, oneshot::Sender<PendingResult>>>>;

pub struct Client {
    sequence: AtomicU16,
    writer: mpsc::UnboundedSender<Vec<u8>>,
    pending: PendingMap,
    pushes: broadcast::Sender<Packet>,
    wire_events: broadcast::Sender<WireEvent>,
    connected: watch::Sender<bool>,
    request_timeout: Duration,
    tasks: Mutex<Vec<JoinHandle<()>>>,
}

impl Client {
    pub async fn connect(config: ClientConfig) -> Result<Self, TransportError> {
        let tcp = timeout(
            config.connect_timeout,
            connect_tcp(&config.host, config.port, config.proxy.as_ref()),
        )
        .await
        .map_err(|_| TransportError::ConnectTimeout)?
        .map_err(|error| TransportError::InvalidConfig(error.to_string()))?;
        tcp.set_nodelay(true).ok();

        let server_name = ServerName::try_from(config.host.clone())
            .map_err(|error| TransportError::InvalidConfig(error.to_string()))?;
        let connector = TlsConnector::from(build_tls_config(config.insecure_tls)?);
        let stream = timeout(config.connect_timeout, connector.connect(server_name, tcp))
            .await
            .map_err(|_| TransportError::ConnectTimeout)?
            .map_err(|error| TransportError::Tls(error.to_string()))?;
        let (mut reader, mut writer) = tokio::io::split(stream);

        let (write_tx, mut write_rx) = mpsc::unbounded_channel::<Vec<u8>>();
        let (push_tx, _) = broadcast::channel(PUSH_CAPACITY);
        let (wire_tx, _) = broadcast::channel(PUSH_CAPACITY);
        let (connected_tx, _) = watch::channel(true);
        let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));

        let writer_connected = connected_tx.clone();
        let writer_pending = pending.clone();
        let writer_task = tokio::spawn(async move {
            while let Some(bytes) = write_rx.recv().await {
                if writer.write_all(&bytes).await.is_err() || writer.flush().await.is_err() {
                    break;
                }
            }
            writer_connected.send_replace(false);
            fail_pending(&writer_pending);
        });

        let reader_connected = connected_tx.clone();
        let reader_pending = pending.clone();
        let reader_pushes = push_tx.clone();
        let reader_wire = wire_tx.clone();
        let reader_task = tokio::spawn(async move {
            let mut decoder = FrameDecoder::default();
            let mut buffer = vec![0u8; READ_BUFFER_SIZE];
            loop {
                let count = match reader.read(&mut buffer).await {
                    Ok(0) | Err(_) => break,
                    Ok(count) => count,
                };
                let frames = match decoder.push(&buffer[..count]) {
                    Ok(frames) => frames,
                    Err(_) => break,
                };
                for frame in frames {
                    let packet = match decode_frame_packet(frame) {
                        Ok(packet) => packet,
                        Err(_) => continue,
                    };
                    let _ = reader_wire.send(WireEvent {
                        direction: Direction::In,
                        packet: packet.clone(),
                    });
                    let response = if packet.command == command::PUSH {
                        None
                    } else {
                        reader_pending.lock().unwrap().remove(&packet.sequence)
                    };
                    if let Some(sender) = response {
                        let _ = sender.send(Ok(packet));
                    } else {
                        let _ = reader_pushes.send(packet);
                    }
                }
            }
            reader_connected.send_replace(false);
            fail_pending(&reader_pending);
        });

        Ok(Self {
            sequence: AtomicU16::new(0),
            writer: write_tx,
            pending,
            pushes: push_tx,
            wire_events: wire_tx,
            connected: connected_tx,
            request_timeout: config.request_timeout,
            tasks: Mutex::new(vec![writer_task, reader_task]),
        })
    }

    pub async fn request(&self, opcode: u16, payload: &[u8]) -> Result<Packet, TransportError> {
        if !self.is_connected() {
            return Err(TransportError::Closed);
        }
        let sequence = self.next_sequence();
        let bytes = encode_request(opcode, sequence, payload)?;
        let (sender, receiver) = oneshot::channel();
        self.pending.lock().unwrap().insert(sequence, sender);
        let _ = self.wire_events.send(WireEvent {
            direction: Direction::Out,
            packet: Packet {
                version: crate::PROTOCOL_VERSION,
                command: command::REQUEST,
                sequence,
                opcode,
                payload: payload.to_vec(),
            },
        });
        if self.writer.send(bytes).is_err() {
            self.pending.lock().unwrap().remove(&sequence);
            return Err(TransportError::Closed);
        }
        match timeout(self.request_timeout, receiver).await {
            Ok(Ok(result)) => result,
            Ok(Err(_)) => Err(TransportError::Closed),
            Err(_) => {
                self.pending.lock().unwrap().remove(&sequence);
                Err(TransportError::Timeout)
            }
        }
    }

    pub fn send(&self, opcode: u16, payload: &[u8]) -> Result<u16, TransportError> {
        if !self.is_connected() {
            return Err(TransportError::Closed);
        }
        let sequence = self.next_sequence();
        let _ = self.wire_events.send(WireEvent {
            direction: Direction::Out,
            packet: Packet {
                version: crate::PROTOCOL_VERSION,
                command: command::REQUEST,
                sequence,
                opcode,
                payload: payload.to_vec(),
            },
        });
        self.writer
            .send(encode_request(opcode, sequence, payload)?)
            .map_err(|_| TransportError::Closed)?;
        Ok(sequence)
    }

    pub fn subscribe(&self) -> broadcast::Receiver<Packet> {
        self.pushes.subscribe()
    }

    pub fn connected(&self) -> watch::Receiver<bool> {
        self.connected.subscribe()
    }

    pub fn wire_events(&self) -> broadcast::Receiver<WireEvent> {
        self.wire_events.subscribe()
    }

    pub fn is_connected(&self) -> bool {
        *self.connected.borrow()
    }

    pub fn close(&self) {
        self.connected.send_replace(false);
        fail_pending(&self.pending);
        for task in self.tasks.lock().unwrap().drain(..) {
            task.abort();
        }
    }

    fn next_sequence(&self) -> u16 {
        self.sequence
            .fetch_add(1, Ordering::Relaxed)
            .wrapping_add(1)
    }
}

impl Drop for Client {
    fn drop(&mut self) {
        self.close();
    }
}

pub fn set_trust_mincifry_ca(enabled: bool) {
    TRUST_MINCIFRY_CA.store(enabled, Ordering::Relaxed);
}

pub fn trust_mincifry_ca() -> bool {
    TRUST_MINCIFRY_CA.load(Ordering::Relaxed)
}

fn fail_pending(pending: &PendingMap) {
    for (_, sender) in pending.lock().unwrap().drain() {
        let _ = sender.send(Err(TransportError::Closed));
    }
}

pub(crate) fn build_tls_config(insecure: bool) -> Result<Arc<RustlsConfig>, TransportError> {
    let provider = Arc::new(ring::default_provider());
    let builder = RustlsConfig::builder_with_provider(provider.clone())
        .with_safe_default_protocol_versions()
        .map_err(|error| TransportError::Tls(error.to_string()))?;
    let config = if insecure {
        builder
            .dangerous()
            .with_custom_certificate_verifier(Arc::new(AcceptAnyCertificate(provider)))
            .with_no_client_auth()
    } else {
        builder
            .with_root_certificates(root_store()?)
            .with_no_client_auth()
    };
    Ok(Arc::new(config))
}

fn root_store() -> Result<RootCertStore, TransportError> {
    let mut roots = RootCertStore {
        roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
    };
    if trust_mincifry_ca() {
        for certificate in rustls_pemfile::certs(&mut MINCIFRY_CA_PEM.as_bytes()) {
            roots
                .add(certificate.map_err(|error| TransportError::Tls(error.to_string()))?)
                .map_err(|error| TransportError::Tls(error.to_string()))?;
        }
    }
    Ok(roots)
}

#[derive(Debug)]
struct AcceptAnyCertificate(Arc<CryptoProvider>);

impl ServerCertVerifier for AcceptAnyCertificate {
    fn verify_server_cert(
        &self,
        _end_entity: &CertificateDer<'_>,
        _intermediates: &[CertificateDer<'_>],
        _server_name: &ServerName<'_>,
        _ocsp_response: &[u8],
        _now: UnixTime,
    ) -> Result<ServerCertVerified, rustls::Error> {
        Ok(ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        message: &[u8],
        certificate: &CertificateDer<'_>,
        signature: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        verify_tls12_signature(
            message,
            certificate,
            signature,
            &self.0.signature_verification_algorithms,
        )
    }

    fn verify_tls13_signature(
        &self,
        message: &[u8],
        certificate: &CertificateDer<'_>,
        signature: &DigitallySignedStruct,
    ) -> Result<HandshakeSignatureValid, rustls::Error> {
        verify_tls13_signature(
            message,
            certificate,
            signature,
            &self.0.signature_verification_algorithms,
        )
    }

    fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
        self.0.signature_verification_algorithms.supported_schemes()
    }
}

#[cfg(test)]
mod tests {
    use super::{root_store, set_trust_mincifry_ca, trust_mincifry_ca, MINCIFRY_CA_PEM};

    #[test]
    fn custom_ca_bundle_contains_two_certificates() {
        let certificates = rustls_pemfile::certs(&mut MINCIFRY_CA_PEM.as_bytes())
            .collect::<Result<Vec<_>, _>>()
            .unwrap();
        assert_eq!(certificates.len(), 2);
    }

    #[test]
    fn custom_ca_is_opt_in() {
        set_trust_mincifry_ca(false);
        let regular = root_store().unwrap().len();
        set_trust_mincifry_ca(true);
        assert!(trust_mincifry_ca());
        assert_eq!(root_store().unwrap().len(), regular + 2);
        set_trust_mincifry_ca(false);
    }
}
