use tracing::*;

use crate::frame::{
    self, ChallengeData, ClientFinishData, ClientHelloData, ClientQueryData, FrameBody,
    GoodbyeData, ServerDumpData, ServerHelloData,
};
use crate::traits::AuthConfig;
use crate::{auth, peer, topic};

/// External action we should take after processing a message.
#[derive(Clone, Debug)]
pub enum Action {
    /// Sends a frame to the other side.
    SendFrame(frame::FrameBody),

    /// Completes the handshake.
    Ready(ReadyData),

    /// Aborts the session, sending a goodbye message.
    Abort(GoodbyeData),

    /// Exits the handshake, with the received goodbye message.
    Exit(GoodbyeData),

    /// Sets the agent of the remote peer.
    SetRemoteAgent(String),

    /// Sets the identity of the remote peer.
    SetRemoteIdent(peer::Identity),

    /// Connection closed gracefully without event.
    HappyClose,
}

/// Generic enum describing the kind of state the handshake is in.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum StateKind {
    WaitHello,
    WaitFinish,
    Ready,
    Fail,
}

/// Client handshake state machine.
#[derive(Clone)]
pub enum ClientState<A: AuthConfig> {
    WaitServerHello {
        auth: A,
        protocol: topic::Topic,
        sent_chal: Option<ChallengeData>,
    },
    Ready {
        auth: A,
    },
    Fail,
}

impl<A: AuthConfig> ClientState<A> {
    pub fn auth(&self) -> Option<&A> {
        match self {
            Self::WaitServerHello { auth, .. } => Some(auth),
            Self::Ready { auth } => Some(auth),
            _ => None,
        }
    }

    pub fn kind(&self) -> StateKind {
        match self {
            Self::WaitServerHello { .. } => StateKind::WaitHello,
            Self::Ready { .. } => StateKind::Ready,
            Self::Fail => StateKind::Fail,
        }
    }
}

/// Server handshake state machine.
#[derive(Clone, Debug)]
pub enum ServerState<A: AuthConfig> {
    WaitClientHello {
        auth: A,
    },
    WaitClientFinish {
        auth: A,
        protocol: topic::Topic,
        client_chal: Option<ChallengeData>,
        server_chal: Option<ChallengeData>,
    },
    Ready {
        auth: A,
    },
    Fail,
}

impl<A: AuthConfig> ServerState<A> {
    /// Creates a new instance as though a connection has been established and
    /// we haven't received any frames yet.
    pub fn init(auth: A) -> Self {
        Self::WaitClientHello { auth }
    }

    pub fn auth(&self) -> Option<&A> {
        match self {
            Self::WaitClientHello { auth, .. } => Some(auth),
            Self::WaitClientFinish { auth, .. } => Some(auth),
            Self::Ready { auth } => Some(auth),
            _ => None,
        }
    }

    pub fn kind(&self) -> StateKind {
        match self {
            Self::WaitClientHello { .. } => StateKind::WaitHello,
            Self::WaitClientFinish { .. } => StateKind::WaitFinish,
            Self::Ready { .. } => StateKind::Ready,
            Self::Fail => StateKind::Fail,
        }
    }
}

#[derive(Clone, Debug)]
pub struct ReadyData {
    protocol: topic::Topic,
}

impl ReadyData {
    pub fn new(protocol: topic::Topic) -> Self {
        Self { protocol }
    }

    pub fn protocol(&self) -> topic::Topic {
        self.protocol
    }
}

#[derive(Clone, Debug)]
pub struct ClientMeta {
    agent: String,
}

impl ClientMeta {
    pub fn new(agent: String) -> Self {
        Self { agent }
    }

    pub fn create_query_frame(&self) -> FrameBody {
        FrameBody::ClientQuery(ClientQueryData::new(self.agent.clone()))
    }
}

#[derive(Clone, Debug)]
pub struct ServerMeta {
    agent: String,
    protocols: Vec<topic::Topic>,
}

impl ServerMeta {
    pub fn add_protocol(&mut self, proto: topic::Topic) {
        self.protocols.push(proto);
    }
}

impl ServerMeta {
    pub fn new(agent: String, protocols: Vec<topic::Topic>) -> Self {
        Self { agent, protocols }
    }
}

/// Inits a client session using the parameters we want to use.
pub fn exec_client_init<A: AuthConfig>(
    protocol: topic::Topic,
    auth: A,
    meta: &ClientMeta,
) -> (ClientState<A>, Vec<Action>) {
    // If we should sign in this auth setup, then generate a challenge here.
    let sent_chal = if auth.get_intent().should_exchange_chals() {
        let chal = auth::gen_challenge();
        Some(chal)
    } else {
        None
    };

    let hello = ClientHelloData::new(meta.agent.clone(), protocol, sent_chal);
    let frame = FrameBody::ClientHello(hello);
    (
        ClientState::WaitServerHello {
            auth,
            protocol,
            sent_chal,
        },
        vec![Action::SendFrame(frame)],
    )
}

/// Performs client handshake state transition on receiving a new frame.
pub fn exec_client_inp<A: AuthConfig>(
    state: &ClientState<A>,
    inp: &FrameBody,
    _meta: &ClientMeta,
) -> (ClientState<A>, Vec<Action>) {
    use ClientState::*;
    match (state, inp) {
        (
            WaitServerHello {
                auth,
                protocol,
                sent_chal,
            },
            FrameBody::ServerHello(shello),
        ) => {
            let mut actions = Vec::new();
            actions.push(Action::SetRemoteAgent(shello.agent().to_owned()));

            let intent = auth.get_intent();
            let should_sign = intent.should_sign(frame::Side::Client);
            let exp_resp = intent.should_sign(frame::Side::Server);
            let client_chal = sent_chal.as_ref();
            let server_chal = shello.challenge();
            let server_resp = shello.response();

            let mut sent_resp = None;

            // Big complicated logic tree to see what we should do about auth.
            if intent.should_exchange_chals() {
                let (Some(client_chal), Some(server_chal)) = (client_chal, server_chal) else {
                    // Missing expected challenge.
                    warn!("missing expected challenge");
                    let goodbye = GoodbyeData::new(-104, "expected chal".to_string());
                    return (ClientState::Fail, vec![Action::Abort(goodbye)]);
                };

                if exp_resp {
                    let Some(server_resp) = server_resp else {
                        // Missing expected response to challenge.
                        warn!("missing expected response to challenge");
                        let goodbye = GoodbyeData::new(-103, "expected chal resp".to_string());
                        return (ClientState::Fail, vec![Action::Abort(goodbye)]);
                    };

                    // Ok, let's check it.
                    match auth.verify_response(
                        client_chal,
                        server_chal,
                        frame::Side::Server,
                        server_resp,
                    ) {
                        Ok(Some(id)) => {
                            trace!(%id, "verified identity");
                            actions.push(Action::SetRemoteIdent(id));
                        }
                        Ok(None) => { /* do nothing */ }
                        Err(e) => {
                            // TODO do something with the error?
                            warn!(err = %e, "server sent bad challenge response");
                            let goodbye = GoodbyeData::new(-102, "auth error".to_string());
                            return (ClientState::Fail, vec![Action::Abort(goodbye)]);
                        }
                    }
                }

                // And we should still sign ourselves.
                if should_sign {
                    match auth.sign_challenge(client_chal, server_chal, frame::Side::Client) {
                        Ok(Some(resp)) => {
                            sent_resp = Some(resp);
                        }
                        Ok(None) => {
                            warn!("intent said we should sign but generated no response");
                        }
                        Err(e) => {
                            error!(err = %e, "could not generate response");
                            let goodbye = GoodbyeData::new(-101, "internal error".to_string());
                            return (ClientState::Fail, vec![Action::Abort(goodbye)]);
                        }
                    }
                }
            }

            let cfinish = ClientFinishData::new(sent_resp);
            let frame = FrameBody::ClientFinish(cfinish);
            let rd = ReadyData::new(*protocol);
            actions.push(Action::SendFrame(frame));
            actions.push(Action::Ready(rd));
            (Ready { auth: auth.clone() }, actions)
        }

        (WaitServerHello { .. }, FrameBody::Goodbye(gd)) => {
            // We sent a protocol the server didn't support.
            (ClientState::Fail, vec![Action::Exit(gd.clone())])
        }

        _ => {
            let goodbye = GoodbyeData::new(-101, "idk what you're doing".to_string());
            (state.clone(), vec![Action::Abort(goodbye)])
        }
    }
}

/// Performs server handshake state transition on receiving a new frame.
pub fn exec_server_inp<A: AuthConfig>(
    state: &ServerState<A>,
    inp: &FrameBody,
    meta: &ServerMeta,
) -> (ServerState<A>, Vec<Action>) {
    use ServerState::*;
    match (state, inp) {
        (WaitClientHello { .. }, FrameBody::ClientQuery(_query)) => {
            let dump = ServerDumpData::new(meta.agent.clone(), meta.protocols.clone());
            let frame = FrameBody::ServerDump(dump);
            (state.clone(), vec![Action::SendFrame(frame)])
        }

        (WaitClientHello { auth }, FrameBody::ClientHello(chello)) => {
            let mut actions = Vec::new();
            actions.push(Action::SetRemoteAgent(chello.agent().to_owned()));

            let protocol = chello.protocol();

            // Check to see they picked a good protocol.
            if !meta.protocols.contains(&protocol) {
                debug!(%protocol, "connection wanted unsupported protocol");
                let goodbye = GoodbyeData::new(-100, "unsupported protocol".to_string());
                return (ServerState::Fail, vec![Action::Abort(goodbye)]);
            }

            // Check to see if we should send across a challenge.
            let intent = auth.get_intent();
            let mut sending_chal = None;
            let mut sending_resp = None;

            if intent.should_exchange_chals() {
                let Some(client_chal) = chello.challenge() else {
                    // Expected the client to send a challenge.
                    let goodbye = GoodbyeData::new(-104, "expected chal".to_string());
                    return (ServerState::Fail, vec![Action::Abort(goodbye)]);
                };

                let server_chal = auth::gen_challenge();
                sending_chal = Some(server_chal);

                if intent.should_sign(frame::Side::Server) {
                    match auth.sign_challenge(client_chal, &server_chal, frame::Side::Server) {
                        Ok(Some(resp)) => {
                            sending_resp = Some(resp);
                        }
                        Ok(None) => {
                            warn!("intent said we should sign but generated no response");
                        }
                        Err(e) => {
                            error!(err = %e, "could not generate response");
                            let goodbye = GoodbyeData::new(-101, "internal error".to_string());
                            return (ServerState::Fail, vec![Action::Abort(goodbye)]);
                        }
                    }
                }
            }

            let shello = ServerHelloData::new(meta.agent.clone(), sending_resp, sending_chal);
            let frame = FrameBody::ServerHello(shello);
            actions.push(Action::SendFrame(frame));

            (
                ServerState::WaitClientFinish {
                    auth: auth.clone(),
                    protocol,
                    client_chal: chello.challenge().cloned(),
                    server_chal: sending_chal,
                },
                actions,
            )
        }

        (
            WaitClientFinish {
                auth,
                protocol,
                client_chal,
                server_chal,
            },
            FrameBody::ClientFinish(cfinish),
        ) => {
            let mut actions = Vec::new();

            // If we expect the client to send a response, go and extract it.
            let intent = auth.get_intent();
            if intent.should_exchange_chals() && intent.should_sign(frame::Side::Client) {
                let (Some(client_chal), Some(server_chal)) = (client_chal, server_chal) else {
                    error!("in an inconsistent state!  didn't have local expected challenges!");
                    let goodbye = GoodbyeData::new(-101, "internal error".to_string());
                    return (ServerState::Fail, vec![Action::Abort(goodbye)]);
                };

                let Some(client_resp) = cfinish.response() else {
                    // Missing expected response to challenge.
                    warn!("missing expected response to challenge");
                    let goodbye = GoodbyeData::new(-103, "expected chal resp".to_string());
                    return (ServerState::Fail, vec![Action::Abort(goodbye)]);
                };

                match auth.verify_response(
                    client_chal,
                    server_chal,
                    frame::Side::Client,
                    client_resp,
                ) {
                    Ok(Some(id)) => {
                        trace!(%id, "verified identity");
                        actions.push(Action::SetRemoteIdent(id))
                    }
                    Ok(None) => {
                        // what should we do here?  pass?
                    }
                    Err(e) => {
                        // Invalid credential.
                        warn!(err = %e, "new client sent bad challenge response");
                        let goodbye = GoodbyeData::new(-102, "auth error".to_string());
                        return (ServerState::Fail, vec![Action::Abort(goodbye)]);
                    }
                }
            }

            // TODO check identification
            let rd = ReadyData::new(*protocol);
            actions.push(Action::Ready(rd));
            (ServerState::Ready { auth: auth.clone() }, actions)
        }

        _ => {
            let goodbye = GoodbyeData::new(-101, "idk what you're doing".to_string());
            (ServerState::Fail, vec![Action::Abort(goodbye)])
        }
    }
}