use crate::native::{ForwardingRequest, NativeClientHello, NativeServerHello, NativeUserAuth};
#[cfg(unix)]
use crate::native::{
    is_supported_user_key_algorithm, parse_ssh_ed25519_public_blob, user_auth_transcript,
};
#[cfg(unix)]
use anyhow::{Context, bail};
use anyhow::{Result, anyhow};
#[cfg(unix)]
use std::io::{Read, Write};
#[cfg(unix)]
use std::os::unix::net::UnixStream;
use std::path::Path;

#[cfg(unix)]
const SSH_AGENT_FAILURE: u8 = 5;
#[cfg(unix)]
const SSH2_AGENTC_REQUEST_IDENTITIES: u8 = 11;
#[cfg(unix)]
const SSH2_AGENT_IDENTITIES_ANSWER: u8 = 12;
#[cfg(unix)]
const SSH2_AGENTC_SIGN_REQUEST: u8 = 13;
#[cfg(unix)]
const SSH2_AGENT_SIGN_RESPONSE: u8 = 14;
#[cfg(unix)]
const SSH_AGENT_RSA_SHA2_512: u32 = 4;
#[cfg(unix)]
const MAX_AGENT_PACKET: usize = 256 * 1024;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AgentIdentity {
    pub key_blob: Vec<u8>,
    pub public_key_algorithm: String,
    pub public_key: Vec<u8>,
    pub sign_algorithm: String,
    pub sign_flags: u32,
    pub comment: String,
}

#[cfg(unix)]
pub fn sign_user_auth_with_agent(
    client: &NativeClientHello,
    server: &NativeServerHello,
    requested_forwardings: Vec<ForwardingRequest>,
) -> Result<NativeUserAuth> {
    let sock = std::env::var("SSH_AUTH_SOCK").context("SSH_AUTH_SOCK is not set")?;
    sign_user_auth_with_agent_at(sock, client, server, requested_forwardings)
}

#[cfg(not(unix))]
pub fn sign_user_auth_with_agent(
    _client: &NativeClientHello,
    _server: &NativeServerHello,
    _requested_forwardings: Vec<ForwardingRequest>,
) -> Result<NativeUserAuth> {
    Err(anyhow!(
        "ssh-agent native auth is not supported on this platform yet; use identity_files"
    ))
}

#[cfg(unix)]
pub fn sign_user_auth_with_agent_at(
    socket_path: impl AsRef<Path>,
    client: &NativeClientHello,
    server: &NativeServerHello,
    requested_forwardings: Vec<ForwardingRequest>,
) -> Result<NativeUserAuth> {
    let mut agent = UnixStream::connect(socket_path.as_ref())
        .with_context(|| format!("connect ssh-agent {}", socket_path.as_ref().display()))?;
    let identities = request_supported_identities(&mut agent)?;
    let identity = identities
        .first()
        .ok_or_else(|| anyhow!("ssh-agent has no supported identities"))?;
    let mut auth = NativeUserAuth {
        public_key_algorithm: identity.sign_algorithm.clone(),
        public_key: identity.public_key.clone(),
        signature: Vec::new(),
        requested_forwardings,
    };
    let transcript = user_auth_transcript(client, server, &auth)?;
    auth.signature = sign_with_agent(&mut agent, identity, &transcript)
        .with_context(|| format!("ssh-agent sign with {}", identity.comment))?;
    Ok(auth)
}

#[cfg(not(unix))]
pub fn sign_user_auth_with_agent_at(
    _socket_path: impl AsRef<Path>,
    _client: &NativeClientHello,
    _server: &NativeServerHello,
    _requested_forwardings: Vec<ForwardingRequest>,
) -> Result<NativeUserAuth> {
    Err(anyhow!(
        "ssh-agent native auth is not supported on this platform yet; use identity_files"
    ))
}

#[cfg(unix)]
fn request_supported_identities(agent: &mut UnixStream) -> Result<Vec<AgentIdentity>> {
    write_agent_packet(agent, &[SSH2_AGENTC_REQUEST_IDENTITIES])?;
    let payload = read_agent_packet(agent)?;
    let mut cursor = payload.as_slice();
    let response = read_u8(&mut cursor)?;
    if response == SSH_AGENT_FAILURE {
        bail!("ssh-agent refused identity request");
    }
    anyhow::ensure!(
        response == SSH2_AGENT_IDENTITIES_ANSWER,
        "unexpected ssh-agent identity response {response}"
    );
    let count = read_u32(&mut cursor)? as usize;
    anyhow::ensure!(count <= 4096, "ssh-agent returned too many identities");
    let mut identities = Vec::new();
    for _ in 0..count {
        let key_blob = read_ssh_string(&mut cursor)?.to_vec();
        let comment = String::from_utf8_lossy(read_ssh_string(&mut cursor)?).to_string();
        if let Some(identity) = supported_identity(key_blob, comment)? {
            identities.push(identity);
        }
    }
    anyhow::ensure!(cursor.is_empty(), "trailing data in ssh-agent identities");
    Ok(identities)
}

#[cfg(unix)]
fn supported_identity(key_blob: Vec<u8>, comment: String) -> Result<Option<AgentIdentity>> {
    let algorithm = key_blob_algorithm(&key_blob)?;
    if !is_supported_user_key_algorithm(&algorithm) {
        return Ok(None);
    }
    let identity = match algorithm.as_str() {
        "ssh-ed25519" => AgentIdentity {
            public_key_algorithm: algorithm.clone(),
            public_key: parse_ssh_ed25519_public_blob(&key_blob)?.to_vec(),
            sign_algorithm: "ssh-ed25519".to_string(),
            sign_flags: 0,
            key_blob,
            comment,
        },
        "ecdsa-sha2-nistp256" => AgentIdentity {
            public_key_algorithm: algorithm.clone(),
            public_key: key_blob.clone(),
            sign_algorithm: algorithm,
            sign_flags: 0,
            key_blob,
            comment,
        },
        "ssh-rsa" => AgentIdentity {
            public_key_algorithm: algorithm,
            public_key: key_blob.clone(),
            sign_algorithm: "rsa-sha2-512".to_string(),
            sign_flags: SSH_AGENT_RSA_SHA2_512,
            key_blob,
            comment,
        },
        _ => return Ok(None),
    };
    Ok(Some(identity))
}

#[cfg(unix)]
fn sign_with_agent(
    agent: &mut UnixStream,
    identity: &AgentIdentity,
    transcript: &[u8],
) -> Result<Vec<u8>> {
    let mut request = Vec::new();
    request.push(SSH2_AGENTC_SIGN_REQUEST);
    write_ssh_string(&mut request, &identity.key_blob);
    write_ssh_string(&mut request, transcript);
    request.extend_from_slice(&identity.sign_flags.to_be_bytes());
    write_agent_packet(agent, &request)?;

    let payload = read_agent_packet(agent)?;
    let mut cursor = payload.as_slice();
    let response = read_u8(&mut cursor)?;
    if response == SSH_AGENT_FAILURE {
        bail!("ssh-agent refused signature request");
    }
    anyhow::ensure!(
        response == SSH2_AGENT_SIGN_RESPONSE,
        "unexpected ssh-agent sign response {response}"
    );
    let signature_blob = read_ssh_string(&mut cursor)?;
    anyhow::ensure!(
        cursor.is_empty(),
        "trailing data in ssh-agent sign response"
    );

    let mut signature_cursor = signature_blob;
    let algorithm = read_ssh_string(&mut signature_cursor)?;
    let algorithm = String::from_utf8_lossy(algorithm);
    anyhow::ensure!(
        algorithm == identity.sign_algorithm,
        "ssh-agent returned signature algorithm {algorithm}, expected {}",
        identity.sign_algorithm
    );
    let signature = read_ssh_string(&mut signature_cursor)?;
    anyhow::ensure!(
        signature_cursor.is_empty(),
        "trailing data in ssh-agent signature blob"
    );
    Ok(signature.to_vec())
}

#[cfg(unix)]
fn read_agent_packet(stream: &mut UnixStream) -> Result<Vec<u8>> {
    let mut len = [0u8; 4];
    stream
        .read_exact(&mut len)
        .context("read ssh-agent packet length")?;
    let len = u32::from_be_bytes(len) as usize;
    anyhow::ensure!(len <= MAX_AGENT_PACKET, "ssh-agent packet too large");
    let mut payload = vec![0u8; len];
    stream
        .read_exact(&mut payload)
        .context("read ssh-agent packet body")?;
    Ok(payload)
}

#[cfg(unix)]
fn write_agent_packet(stream: &mut UnixStream, payload: &[u8]) -> Result<()> {
    anyhow::ensure!(
        payload.len() <= MAX_AGENT_PACKET,
        "ssh-agent request too large"
    );
    stream.write_all(&(payload.len() as u32).to_be_bytes())?;
    stream.write_all(payload)?;
    Ok(())
}

#[cfg(unix)]
fn read_u8(cursor: &mut &[u8]) -> Result<u8> {
    anyhow::ensure!(!cursor.is_empty(), "truncated u8");
    let value = cursor[0];
    *cursor = &cursor[1..];
    Ok(value)
}

#[cfg(unix)]
fn read_u32(cursor: &mut &[u8]) -> Result<u32> {
    anyhow::ensure!(cursor.len() >= 4, "truncated u32");
    let value = u32::from_be_bytes(cursor[..4].try_into().unwrap());
    *cursor = &cursor[4..];
    Ok(value)
}

#[cfg(unix)]
fn read_ssh_string<'a>(cursor: &mut &'a [u8]) -> Result<&'a [u8]> {
    let len = read_u32(cursor)? as usize;
    anyhow::ensure!(cursor.len() >= len, "truncated SSH string");
    let (value, rest) = cursor.split_at(len);
    *cursor = rest;
    Ok(value)
}

#[cfg(unix)]
fn write_ssh_string(out: &mut Vec<u8>, value: &[u8]) {
    out.extend_from_slice(&(value.len() as u32).to_be_bytes());
    out.extend_from_slice(value);
}

#[cfg(unix)]
fn key_blob_algorithm(blob: &[u8]) -> Result<String> {
    let mut cursor = blob;
    let algorithm = read_ssh_string(&mut cursor)?;
    Ok(String::from_utf8_lossy(algorithm).to_string())
}

#[cfg(all(test, unix))]
mod tests {
    use super::*;
    use crate::native::{
        host_public_key, parse_authorized_keys, sign_server_hello, ssh_ed25519_public_blob,
        verify_native_user_auth,
    };
    use base64::Engine;
    use base64::engine::general_purpose::STANDARD;
    use ed25519_dalek::{Signer, SigningKey, VerifyingKey};
    use std::os::unix::net::UnixListener;
    use std::thread;

    #[test]
    fn agent_ed25519_signature_authenticates_native_user() {
        let dir = tempfile::tempdir().unwrap();
        let socket_path = dir.path().join("agent.sock");
        let user_signing = SigningKey::from_bytes(&[51u8; 32]);
        let user_public = VerifyingKey::from(&user_signing).to_bytes();
        let key_blob = ssh_ed25519_public_blob(&user_public);
        let listener = UnixListener::bind(&socket_path).unwrap();
        thread::spawn(move || fake_agent(listener, key_blob, user_signing));

        let host_signing = SigningKey::from_bytes(&[3u8; 32]);
        let client = test_client_hello();
        let mut server = test_server_hello(&host_signing);
        sign_server_hello(&host_signing, &client, &mut server).unwrap();
        let auth =
            sign_user_auth_with_agent_at(&socket_path, &client, &server, Vec::new()).unwrap();
        let authorized_line = format!(
            "ssh-ed25519 {} agent-key",
            STANDARD.encode(ssh_ed25519_public_blob(&user_public))
        );
        let authorized = parse_authorized_keys(&authorized_line).unwrap();

        verify_native_user_auth(&client, &server, &auth, &authorized, None).unwrap();
    }

    fn fake_agent(listener: UnixListener, key_blob: Vec<u8>, signing: SigningKey) {
        let (mut stream, _) = listener.accept().unwrap();
        let payload = read_agent_packet(&mut stream).unwrap();
        assert_eq!(payload, vec![SSH2_AGENTC_REQUEST_IDENTITIES]);
        let mut response = Vec::new();
        response.push(SSH2_AGENT_IDENTITIES_ANSWER);
        response.extend_from_slice(&1u32.to_be_bytes());
        write_ssh_string(&mut response, &key_blob);
        write_ssh_string(&mut response, b"test-agent-key");
        write_agent_packet(&mut stream, &response).unwrap();

        let payload = read_agent_packet(&mut stream).unwrap();
        let mut cursor = payload.as_slice();
        assert_eq!(read_u8(&mut cursor).unwrap(), SSH2_AGENTC_SIGN_REQUEST);
        assert_eq!(read_ssh_string(&mut cursor).unwrap(), key_blob.as_slice());
        let transcript = read_ssh_string(&mut cursor).unwrap();
        assert_eq!(read_u32(&mut cursor).unwrap(), 0);
        assert!(cursor.is_empty());
        let signature = signing.sign(transcript).to_bytes();
        let mut signature_blob = Vec::new();
        write_ssh_string(&mut signature_blob, b"ssh-ed25519");
        write_ssh_string(&mut signature_blob, &signature);
        let mut response = Vec::new();
        response.push(SSH2_AGENT_SIGN_RESPONSE);
        write_ssh_string(&mut response, &signature_blob);
        write_agent_packet(&mut stream, &response).unwrap();
    }

    fn test_client_hello() -> NativeClientHello {
        NativeClientHello {
            protocol_version: crate::native::NATIVE_PROTOCOL_VERSION,
            client_random: [1u8; 32],
            client_ephemeral_public: [2u8; 32],
            requested_host: "homelab".to_string(),
            requested_user: "alice".to_string(),
            requested_session: "term".to_string(),
            requested_mode: "read-write".to_string(),
            terminal_size: (80, 24),
            supported_aead: vec!["chacha20poly1305".to_string()],
            supported_user_key_algorithms: vec!["ssh-ed25519".to_string()],
            cached_host_key_fingerprint: None,
            attach_ticket_envelope: None,
            requested_env: Vec::new(),
        }
    }

    fn test_server_hello(signing: &SigningKey) -> NativeServerHello {
        NativeServerHello {
            protocol_version: crate::native::NATIVE_PROTOCOL_VERSION,
            server_random: [4u8; 32],
            server_ephemeral_public: [5u8; 32],
            host_key: host_public_key(signing),
            chosen_aead: "chacha20poly1305".to_string(),
            server_key_epoch: 1,
            auth_challenge: [6u8; 32],
            rate_limit_remaining: Some(30),
            host_signature: Vec::new(),
        }
    }
}