use crate::api::ConnectionSettings;
use crate::{opamp::*, state::*};
use async_trait::async_trait;
use futures_util::SinkExt;
use futures_util::StreamExt;
use prost::Message as ProstMessage;
use std::any::Any;
use sysinfo::{System, SystemExt};
use tokio::net::TcpStream;
use tokio::runtime;
use tokio_tungstenite::{
    connect_async, tungstenite::protocol::Message, MaybeTlsStream, WebSocketStream,
};

pub struct WsClient {
    settings: ConnectionSettings,
    address: url::Url,
    backoff: u32,
    seqno: u64,
    agent_state: Option<AgentToServer>,
    stream: Option<WebSocketStream<MaybeTlsStream<TcpStream>>>,
}

impl WsClient {
    pub fn new(settings: ConnectionSettings) -> WsClient {
        let path = settings.server_endpoint.clone() + settings.listen_path.as_str();
        let address = url::Url::parse(&path).unwrap();

        WsClient {
            settings,
            address,
            backoff: 0,
            seqno: 0,
            agent_state: None,
            stream: None,
        }
    }

    pub async fn send(&mut self, buf: Message) -> Result<(), OpAMPClientError> {
        let (mut s, _) = self.stream.as_mut().unwrap().split();
        match s.send(buf).await {
            Ok(_) => Ok(()),
            Err(e) => Err(OpAMPClientError::new(
                line!(),
                format!("Websocket send error {}", e).as_str(),
            )),
        }
    }

    pub async fn receive(&mut self) -> Result<Option<Message>, OpAMPClientError> {
        let (_, mut r) = self.stream.as_mut().unwrap().split();
        match r.next().await {
            Some(Ok(message)) => return Ok(Some(message)),
            Some(Err(e)) => {
                return Err(OpAMPClientError {
                    code: line!(),
                    details: format!("WebSocket receive error {}", e).to_string(),
                })
            }
            None => return Ok(None),
        }
    }

    pub fn get_time_nanos(&self) -> u128 {
        let start = std::time::SystemTime::now();
        let since_the_epoch = start
            .duration_since(std::time::UNIX_EPOCH)
            .expect("Time went backwards");
        since_the_epoch.as_nanos()
    }

    fn set_health(&mut self, healthy: bool) {
        // We're now polling. Set our state to healthy
        if let Some(state) = &mut self.agent_state {
            if let Some(health) = &mut state.health {
                health.healthy = healthy;
            }
        }
    }

    pub async fn report_status(&mut self) -> Result<(), OpAMPClientError> {
        // Populate an initial state if it doesnt yet exist
        if let None = self.agent_state {
            // ---// Agent Description //---
            let mut ad = AgentDescription {
                identifying_attributes: vec![],
                non_identifying_attributes: vec![],
            };
            // Identifying attributes
            ad.identifying_attributes.push(KeyValue {
                key: "service.name".to_string(),
                value: Some(AnyValue {
                    value: Some(any_value::Value::StringValue(self.settings.name.clone())),
                }),
            });
            ad.identifying_attributes.push(KeyValue {
                key: "service.version".to_string(),
                value: Some(AnyValue {
                    value: Some(any_value::Value::StringValue(self.settings.version.clone())),
                }),
            });
            ad.identifying_attributes.push(KeyValue {
                key: "service.instance.id".to_string(),
                value: Some(AnyValue {
                    value: Some(any_value::Value::StringValue(
                        self.settings.instance_id.clone(),
                    )),
                }),
            });
            // Non-identifying attributes
            let sys = System::new_all();
            ad.non_identifying_attributes.push(KeyValue {
                key: "os.type".to_string(),
                value: Some(AnyValue {
                    value: Some(any_value::Value::StringValue(
                        std::env::consts::OS.to_string(),
                    )),
                }),
            });
            ad.non_identifying_attributes.push(KeyValue {
                key: "os.version".to_string(),
                value: Some(AnyValue {
                    value: Some(any_value::Value::StringValue(sys.kernel_version().unwrap())),
                }),
            });
            ad.non_identifying_attributes.push(KeyValue {
                key: "host.name".to_string(),
                value: Some(AnyValue {
                    value: Some(any_value::Value::StringValue(sys.host_name().unwrap())),
                }),
            });
            let agent_description = Some(ad);

            // ---// Health //---
            let health = Some(AgentHealth {
                healthy: false, // Note: Initial state is always false
                start_time_unix_nano: self.get_time_nanos() as u64,
                last_error: "".to_string(),
            });

            // ---// Effective Configuration //---
            let effective_config = Some(EffectiveConfig { config_map: None });

            // ---// Remote Configuration Status //---
            let remote_config_status = Some(RemoteConfigStatus {
                last_remote_config_hash: vec![],
                status: RemoteConfigStatuses::Unset.into(),
                error_message: "".to_string(),
            });

            // ---// Package Status //---
            let package_statuses = Some(PackageStatuses {
                packages: std::collections::HashMap::new(),
                server_provided_all_packages_hash: vec![],
                error_message: "".to_string(),
            });

            self.seqno += 1;
            // Set a few defaults (TODO: Revisit)
            let cap = AgentCapabilities::ReportsHealth as u64
                & AgentCapabilities::ReportsStatus as u64
                & AgentCapabilities::AcceptsPackages as u64
                & AgentCapabilities::ReportsPackageStatuses as u64;
            self.agent_state = Some(AgentToServer {
                instance_uid: self.settings.instance_id.clone(),
                sequence_num: self.seqno,
                agent_description,
                capabilities: cap,
                health,
                effective_config,
                remote_config_status,
                package_statuses,
                agent_disconnect: None,
                flags: 789,
            });
        }

        if let Some(state_message) = &self.agent_state {
            let mut buf = Vec::new();
            state_message.encode(&mut buf).unwrap();
            self.send(Message::Binary(buf)).await.unwrap();
        }

        Ok(())
    }
}

#[async_trait]
impl Client for WsClient {
    async fn connect(&mut self) -> Result<(), OpAMPClientError> {
        self.stream = match connect_async(self.address.clone()).await {
            Ok(s) => {
                let (strm, _) = s;
                Some(strm)
            }
            Err(e) => {
                return Err(OpAMPClientError {
                    code: line!(),
                    details: format!("Failed to connect to websocket: {}", e),
                })
            }
        };

        log::info!("Websocket connection to server successful");

        Ok(())
    }

    async fn send(&mut self, buf: &ClientBuffer) -> Result<(), OpAMPClientError> {
        let ws_message = Message::Binary(buf.clone());
        self.send(ws_message).await
    }



    fn process(&mut self, current_state: &dyn Any) -> bool {
        let rt = runtime::Builder::new_current_thread().build().unwrap();

        if let Some(_) = current_state.downcast_ref::<Disconnected>() {
            return self.stream.is_none();
        }
        if let Some(_) = current_state.downcast_ref::<Connecting>() {
            match rt.block_on(self.connect()) {
                Ok(_) => return true,
                Err(_) => {
                    // Backoff sleep
                    self.backoff += 1;

                    let connect_retries =
                        std::env::var("AMPERSAND_CONNECT_RETRIES").unwrap_or("10".to_string());
                    if self.backoff > connect_retries.parse::<u32>().unwrap() {
                        log::error!("Failed to connect after {} retries", connect_retries);
                        std::process::exit(-1);
                    }
                    let idle_sec = std::time::Duration::from_secs(2_u64.pow(self.backoff));
                    std::thread::sleep(idle_sec);
                    return false;
                }
            }
        }
        if let Some(_) = current_state.downcast_ref::<Connected>() {
            if let Ok(_) = rt.block_on(self.report_status()) {
                self.set_health(true);
                return true;
            }
            return false;
        }
        if let Some(_) = current_state.downcast_ref::<Polling>() {
            // Check the websocket inbound
            let rt = runtime::Builder::new_current_thread().build().unwrap();
            if let Ok(Some(Message::Binary(bytes))) = rt.block_on(self.receive()) {
                if let Ok(msg) =
                    crate::opamp::ServerToAgent::decode(&mut std::io::Cursor::new(&bytes))
                {
                    if let Some(_command) = msg.command {
                        // TODO: invoke onCommand on the api client callback
                    }

                    // Check and report full state
                    if msg.flags & (ServerToAgentFlags::ReportFullState as u64) != 0 {
                        rt.block_on(self.report_status()).unwrap()
                    }

                    if let Some(agent_rc) = msg.remote_config {
                        log::debug!("Received a remote config: {:?}", agent_rc);
                        // TODO: Do something with it
                    }

                    // TODO: Check our agent capabilities if it supports any of these
                    // else ignore them harmlessly
                    if let Some(_connection_settings_offers) = msg.connection_settings {
                        if let Some(_owm_metrics) = _connection_settings_offers.own_metrics {
                            // TODO: Do something with it
                        }
                        if let Some(_owm_traces) = _connection_settings_offers.own_traces {
                            // TODO: Do something with it
                        }
                        if let Some(_owm_logs) = _connection_settings_offers.own_logs {
                            // TODO: Do something with it
                        }
                    }

                    if let Some(_packages_available) = msg.packages_available {
                        // TODO: Initiate a download and unpack of packages
                    }

                    // Check and send messages on publish queue
                    // If something needs to be run locally, return true
                }
            }
            return true;
        }
        if let Some(_) = current_state.downcast_ref::<Executing>() {
            // Run commands in the exec queue

            // Return status as needed

            return true;
        }

        false
    }
}
