use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde_json::{json, Map};
use tokio::sync::broadcast;
use tokio::task::JoinHandle;

use crate::message::{decode_tagged_json, encode_tagged_json, MessageError};
use crate::packet::{command, Packet};
use crate::transport::WireEvent;
use crate::transport::{Client, ClientConfig, TransportError};
use crate::ProxyConfig;

const SESSION_INIT_OPCODE: u16 = 6;
const PING_OPCODE: u16 = 1;
const PUSH_CAPACITY: usize = 256;

#[derive(Clone, Debug)]
pub struct SessionOptions {
    pub host: String,
    pub port: u16,
    pub device_id: String,
    pub instance_id: String,
    pub app_version: String,
    pub build_number: i64,
    pub device_type: String,
    pub os_version: String,
    pub timezone: String,
    pub screen: String,
    pub push_device_type: String,
    pub arch: String,
    pub locale: String,
    pub device_name: String,
    pub device_locale: String,
    pub client_session_id: i64,
    pub ping_interval: Duration,
    pub ping_interactive: bool,
    pub insecure_tls: bool,
    pub proxy: Option<String>,
}

impl SessionOptions {
    pub fn client_config(&self) -> Result<ClientConfig, TransportError> {
        let mut config = ClientConfig::new(self.host.clone(), self.port);
        config.insecure_tls = self.insecure_tls;
        config.proxy = self
            .proxy
            .as_deref()
            .map(ProxyConfig::parse)
            .transpose()
            .map_err(|error| TransportError::InvalidConfig(error.to_string()))?;
        Ok(config)
    }

    pub fn http_user_agent(&self) -> String {
        format!(
            "OKMessages/{} ({}; {}; {})",
            self.app_version, self.os_version, self.device_name, self.screen
        )
    }

    fn handshake_payload(&self) -> Result<Vec<u8>, MessageError> {
        let mut user_agent = Map::new();
        user_agent.insert("deviceType".to_string(), json!(self.device_type));
        user_agent.insert("appVersion".to_string(), json!(self.app_version));
        user_agent.insert("osVersion".to_string(), json!(self.os_version));
        user_agent.insert("timezone".to_string(), json!(self.timezone));
        user_agent.insert("screen".to_string(), json!(self.screen));
        user_agent.insert("pushDeviceType".to_string(), json!(self.push_device_type));
        user_agent.insert("arch".to_string(), json!(self.arch));
        user_agent.insert("locale".to_string(), json!(self.locale));
        user_agent.insert("buildNumber".to_string(), json!(self.build_number));
        user_agent.insert("deviceName".to_string(), json!(self.device_name));
        user_agent.insert("deviceLocale".to_string(), json!(self.device_locale));
        encode_tagged_json(
            &json!({
                "mt_instanceid": self.instance_id,
                "userAgent": user_agent,
                "clientSessionId": self.client_session_id,
                "deviceId": self.device_id,
            })
            .to_string(),
        )
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum SessionState {
    Disconnected = 0,
    Connecting = 1,
    Connected = 2,
    Online = 3,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HandshakeInfo {
    pub calls_seed: Option<i64>,
    pub device_name: Option<String>,
    pub payload: Vec<u8>,
    pub payload_json: String,
}

pub struct Session {
    options: SessionOptions,
    client: Mutex<Option<Arc<Client>>>,
    pushes: broadcast::Sender<Packet>,
    wire_events: broadcast::Sender<WireEvent>,
    state: AtomicU8,
    ping_interactive: Arc<AtomicBool>,
    tasks: Mutex<Vec<JoinHandle<()>>>,
}

impl Session {
    pub fn new(options: SessionOptions) -> Self {
        let (pushes, _) = broadcast::channel(PUSH_CAPACITY);
        let (wire_events, _) = broadcast::channel(PUSH_CAPACITY);
        let ping_interactive = options.ping_interactive;
        Self {
            options,
            client: Mutex::new(None),
            pushes,
            wire_events,
            state: AtomicU8::new(SessionState::Disconnected as u8),
            ping_interactive: Arc::new(AtomicBool::new(ping_interactive)),
            tasks: Mutex::new(Vec::new()),
        }
    }

    pub async fn connect(&self) -> Result<HandshakeInfo, TransportError> {
        self.disconnect();
        self.set_state(SessionState::Connecting);
        let config = match self.options.client_config() {
            Ok(config) => config,
            Err(error) => {
                self.set_state(SessionState::Disconnected);
                return Err(error);
            }
        };
        let client = match Client::connect(config).await {
            Ok(client) => Arc::new(client),
            Err(error) => {
                self.set_state(SessionState::Disconnected);
                return Err(error);
            }
        };
        self.set_state(SessionState::Connected);

        let mut incoming_pushes = client.subscribe();
        let mut incoming_wire = client.wire_events();
        let push_output = self.pushes.clone();
        let push_task = tokio::spawn(async move {
            while let Ok(packet) = incoming_pushes.recv().await {
                let _ = push_output.send(packet);
            }
        });
        let wire_output = self.wire_events.clone();
        let wire_task = tokio::spawn(async move {
            while let Ok(event) = incoming_wire.recv().await {
                let _ = wire_output.send(event);
            }
        });

        let handshake = async {
            let handshake_payload = self
                .options
                .handshake_payload()
                .map_err(|error| TransportError::InvalidConfig(error.to_string()))?;
            let response = client
                .request(SESSION_INIT_OPCODE, &handshake_payload)
                .await?;
            if response.command != command::OK {
                return Err(server_error(&response));
            }
            let payload_json = decode_tagged_json(&response.payload)
                .map_err(|error| TransportError::InvalidConfig(error.to_string()))?;
            let value: serde_json::Value = serde_json::from_str(&payload_json)
                .map_err(|error| TransportError::InvalidConfig(error.to_string()))?;
            Ok(HandshakeInfo {
                calls_seed: value.get("callsSeed").and_then(serde_json::Value::as_i64),
                device_name: value
                    .get("device_name")
                    .and_then(serde_json::Value::as_str)
                    .map(str::to_string),
                payload: response.payload,
                payload_json,
            })
        }
        .await;
        let info = match handshake {
            Ok(info) => info,
            Err(error) => {
                push_task.abort();
                wire_task.abort();
                client.close();
                self.set_state(SessionState::Disconnected);
                return Err(error);
            }
        };

        let ping_client = client.clone();
        let ping_interval = self.options.ping_interval;
        let ping_interactive = self.ping_interactive.clone();
        let ping_task = tokio::spawn(async move {
            let mut interval = tokio::time::interval_at(
                tokio::time::Instant::now() + ping_interval,
                ping_interval,
            );
            loop {
                interval.tick().await;
                let payload = ping_payload(ping_interactive.load(Ordering::Relaxed));
                if ping_client.send(PING_OPCODE, &payload).is_err() {
                    break;
                }
            }
        });

        *self.client.lock().unwrap() = Some(client);
        self.tasks
            .lock()
            .unwrap()
            .extend([push_task, wire_task, ping_task]);
        self.set_state(SessionState::Online);
        Ok(info)
    }

    pub async fn request_raw(&self, opcode: u16, payload: &[u8]) -> Result<Packet, TransportError> {
        let client = self.client.lock().unwrap().clone();
        match client {
            Some(client) => client.request(opcode, payload).await,
            None => Err(TransportError::Closed),
        }
    }

    pub async fn request(&self, opcode: u16, payload: &[u8]) -> Result<Packet, TransportError> {
        let response = self.request_raw(opcode, payload).await?;
        if response.command == command::ERROR {
            return Err(server_error(&response));
        }
        Ok(response)
    }

    pub fn send(&self, opcode: u16, payload: &[u8]) -> Result<u16, TransportError> {
        let client = self.client.lock().unwrap().clone();
        match client {
            Some(client) => client.send(opcode, payload),
            None => Err(TransportError::Closed),
        }
    }

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

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

    pub fn set_ping_interactive(&self, interactive: bool) {
        self.ping_interactive.store(interactive, Ordering::Relaxed);
        let payload = ping_payload(interactive);
        let _ = self.send(PING_OPCODE, &payload);
    }

    pub fn ping_interactive(&self) -> bool {
        self.ping_interactive.load(Ordering::Relaxed)
    }

    pub fn state(&self) -> SessionState {
        match self.state.load(Ordering::Relaxed) {
            1 => SessionState::Connecting,
            2 => SessionState::Connected,
            3 => SessionState::Online,
            _ => SessionState::Disconnected,
        }
    }

    pub fn user_agent(&self) -> String {
        self.options.http_user_agent()
    }

    pub fn disconnect(&self) {
        if let Some(client) = self.client.lock().unwrap().take() {
            client.close();
        }
        for task in self.tasks.lock().unwrap().drain(..) {
            task.abort();
        }
        self.set_state(SessionState::Disconnected);
    }

    fn set_state(&self, state: SessionState) {
        self.state.store(state as u8, Ordering::Relaxed);
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        self.disconnect();
    }
}

fn ping_payload(interactive: bool) -> Vec<u8> {
    encode_tagged_json(&json!({"interactive": interactive}).to_string()).unwrap_or_default()
}

fn server_error(packet: &Packet) -> TransportError {
    let value = decode_tagged_json(&packet.payload)
        .ok()
        .and_then(|json| serde_json::from_str::<serde_json::Value>(&json).ok())
        .unwrap_or(serde_json::Value::Null);
    let message = ["localizedMessage", "message", "title"]
        .iter()
        .find_map(|key| value.get(key).and_then(serde_json::Value::as_str))
        .unwrap_or("unknown server error")
        .to_string();
    let key = value
        .get("error")
        .and_then(serde_json::Value::as_str)
        .map(str::to_string);
    TransportError::Server { message, key }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use rcgen::generate_simple_self_signed;
    use rustls::pki_types::PrivateKeyDer;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpListener;
    use tokio_rustls::TlsAcceptor;

    use super::{ping_payload, Session, SessionOptions, SessionState};
    use crate::{command, decode_packet, decode_tagged_json, encode_packet, encode_tagged_json};

    fn options() -> SessionOptions {
        SessionOptions {
            host: "example.invalid".to_string(),
            port: 443,
            device_id: "synthetic-device".to_string(),
            instance_id: "synthetic-instance".to_string(),
            app_version: "1.0.0".to_string(),
            build_number: 1,
            device_type: "TEST".to_string(),
            os_version: "Synthetic OS".to_string(),
            timezone: "Etc/UTC".to_string(),
            screen: "100x100".to_string(),
            push_device_type: "NONE".to_string(),
            arch: "test-arch".to_string(),
            locale: "zz".to_string(),
            device_name: "Synthetic Device".to_string(),
            device_locale: "zz".to_string(),
            client_session_id: 2,
            ping_interval: Duration::from_secs(30),
            ping_interactive: true,
            insecure_tls: false,
            proxy: None,
        }
    }

    #[test]
    fn handshake_contains_required_session_fields() {
        let json = decode_tagged_json(&options().handshake_payload().unwrap()).unwrap();
        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert_eq!(value["deviceId"], "synthetic-device");
        assert_eq!(value["mt_instanceid"], "synthetic-instance");
        assert_eq!(value["userAgent"]["deviceType"], "TEST");
    }

    #[test]
    fn ping_contains_interactive_state() {
        let json = decode_tagged_json(&ping_payload(false)).unwrap();
        let value: serde_json::Value = serde_json::from_str(&json).unwrap();
        assert_eq!(value["interactive"], false);
    }

    #[tokio::test]
    async fn invalid_proxy_restores_disconnected_state() {
        let mut config = options();
        config.proxy = Some("ftp://127.0.0.1:21".to_string());
        let session = Session::new(config);
        assert!(session.connect().await.is_err());
        assert_eq!(session.state(), SessionState::Disconnected);
    }

    #[tokio::test]
    async fn tls_session_exchanges_request_and_push() {
        let certified = generate_simple_self_signed(vec!["localhost".to_string()]).unwrap();
        let server_config = rustls::ServerConfig::builder()
            .with_no_client_auth()
            .with_single_cert(
                vec![certified.cert.der().clone()],
                PrivateKeyDer::Pkcs8(certified.key_pair.serialize_der().into()),
            )
            .unwrap();
        let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
        let port = listener.local_addr().unwrap().port();
        let server = tokio::spawn(async move {
            let (tcp, _) = listener.accept().await.unwrap();
            let mut stream = TlsAcceptor::from(std::sync::Arc::new(server_config))
                .accept(tcp)
                .await
                .unwrap();
            let handshake = read_packet(&mut stream).await;
            assert_eq!(handshake.opcode, 6);
            let handshake_payload = encode_tagged_json(r#"{"callsSeed":17}"#).unwrap();
            stream
                .write_all(
                    &encode_packet(
                        command::OK,
                        handshake.opcode,
                        handshake.sequence,
                        &handshake_payload,
                    )
                    .unwrap(),
                )
                .await
                .unwrap();
            let push_payload = encode_tagged_json(r#"{"kind":"synthetic"}"#).unwrap();
            stream
                .write_all(&encode_packet(command::PUSH, 137, 0, &push_payload).unwrap())
                .await
                .unwrap();
            let request = read_packet(&mut stream).await;
            assert_eq!(request.opcode, 17);
            let response_payload = encode_tagged_json(r#"{"authorized":true}"#).unwrap();
            stream
                .write_all(
                    &encode_packet(
                        command::OK,
                        request.opcode,
                        request.sequence,
                        &response_payload,
                    )
                    .unwrap(),
                )
                .await
                .unwrap();
        });

        let mut config = options();
        config.host = "localhost".to_string();
        config.port = port;
        config.insecure_tls = true;
        let session = Session::new(config);
        let mut pushes = session.pushes();
        let handshake = session.connect().await.unwrap();
        assert_eq!(handshake.calls_seed, Some(17));
        let request_payload = encode_tagged_json(r#"{"mode":{"$bin":"AQID"}}"#).unwrap();
        let response = session.request(17, &request_payload).await.unwrap();
        let decoded: serde_json::Value =
            serde_json::from_str(&decode_tagged_json(&response.payload).unwrap()).unwrap();
        assert_eq!(decoded["authorized"], true);
        let push = tokio::time::timeout(Duration::from_secs(2), pushes.recv())
            .await
            .unwrap()
            .unwrap();
        assert_eq!(push.opcode, 137);
        server.await.unwrap();
    }

    async fn read_packet<S>(stream: &mut S) -> crate::Packet
    where
        S: tokio::io::AsyncRead + Unpin,
    {
        let mut header = [0u8; 10];
        stream.read_exact(&mut header).await.unwrap();
        let packed = u32::from_be_bytes([header[6], header[7], header[8], header[9]]);
        let length = (packed & 0x00ff_ffff) as usize;
        let mut wire = header.to_vec();
        wire.resize(10 + length, 0);
        stream.read_exact(&mut wire[10..]).await.unwrap();
        decode_packet(&wire).unwrap()
    }
}
