use std::path::Path;

use anyhow::{Context, anyhow};
use data_encoding::BASE64;
use der::{Decode as _, Encode as _, asn1::BitStringRef};
use snow::{
  params::NoiseParams,
  resolvers::{CryptoResolver, DefaultResolver},
};
use spki::SubjectPublicKeyInfoRef;

use crate::PkiKind;

#[derive(PartialEq, Clone)]
pub struct SpkiPublicKey(String);

impl From<String> for SpkiPublicKey {
  fn from(value: String) -> Self {
    Self(value)
  }
}

impl std::fmt::Display for SpkiPublicKey {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.write_str(&self.0)
  }
}

impl SpkiPublicKey {
  pub fn as_str(&self) -> &str {
    &self.0
  }

  pub fn as_bytes(&self) -> &[u8] {
    self.0.as_bytes()
  }

  pub fn into_inner(self) -> String {
    self.0
  }

  pub fn as_pem(&self) -> String {
    let public_key = &self.0;
    format!(
      r#"-----BEGIN PUBLIC KEY-----
{public_key}
-----END PUBLIC KEY-----
"#
    )
  }

  pub fn write_pem_sync(
    &self,
    path: impl AsRef<Path>,
  ) -> anyhow::Result<()> {
    let path = path.as_ref();
    tracing::info!("Writing public key to {path:?}");
    mogh_secret_file::write(path, self.as_pem()).with_context(|| {
      format!("Failed to write public key pem to {path:?}")
    })
  }

  pub async fn write_pem_async(
    &self,
    path: impl AsRef<Path>,
  ) -> anyhow::Result<()> {
    let path = path.as_ref();
    tracing::info!("Writing public key to {path:?}");
    mogh_secret_file::write_async(path, self.as_pem())
      .await
      .with_context(|| {
        format!("Failed to write public key pem to {path:?}")
      })
  }

  /// Supports file or hardcoded spec.
  ///
  /// - Direct pass: `public_key = "MCow..."`
  /// - File path: `public_key = "file:/path/to/key.pub"`
  pub fn from_spec(spec: &str) -> anyhow::Result<Self> {
    if let Some(path) = spec.strip_prefix("file:") {
      SpkiPublicKey::from_file(path)
    } else {
      SpkiPublicKey::from_maybe_pem(spec)
    }
  }

  pub fn from_file(path: impl AsRef<Path>) -> anyhow::Result<Self> {
    let path = path.as_ref();
    let contents =
      std::fs::read_to_string(path).with_context(|| {
        format!("Failed to read public key at {path:?}")
      })?;
    Self::from_maybe_pem(&contents)
  }

  /// Accepts pem rfc7468 (openssl) or base64 der (second line of rfc7468 pem).
  pub fn from_maybe_pem(
    public_key_maybe_pem: &str,
  ) -> anyhow::Result<Self> {
    let public_key_der =
      Self::maybe_pem_to_der(public_key_maybe_pem)?;
    Self::from_der(&public_key_der)
  }

  /// Accepts der (not base64)
  pub fn from_der(public_key_der: &[u8]) -> anyhow::Result<Self> {
    let spki = SubjectPublicKeyInfoRef::from_der(public_key_der)
      .map_err(anyhow::Error::msg)
      .context("Invalid public key der")?;
    if spki.algorithm.oid != super::OID_X25519 {
      return Err(anyhow!("Public key is not X25519"));
    }
    Self::from_raw_bytes(spki.subject_public_key.raw_bytes())
  }

  pub fn from_raw_bytes(public_key: &[u8]) -> anyhow::Result<Self> {
    let bs = BitStringRef::new(0, public_key)
      .map_err(anyhow::Error::msg)
      .context("Failed to parse public key bytes into bit string")?;

    let spki = spki::SubjectPublicKeyInfo {
      algorithm: super::algorithm(),
      subject_public_key: bs,
    };

    let mut buf = [0u8; 128];
    let public_key = spki
      .encode_to_slice(&mut buf)
      .map_err(anyhow::Error::msg)
      .context("Failed to write subject public key info into der")?;

    Ok(Self(BASE64.encode(public_key)))
  }

  pub fn maybe_pem_to_raw_bytes(
    public_key_maybe_pem: &str,
  ) -> anyhow::Result<[u8; 32]> {
    let der = Self::maybe_pem_to_der(public_key_maybe_pem)?;
    Self::der_to_raw_bytes(&der)
  }

  pub fn maybe_pem_to_der(
    public_key_maybe_pem: &str,
  ) -> anyhow::Result<Vec<u8>> {
    if public_key_maybe_pem.starts_with("-----BEGIN") {
      let (_label, public_key_der) =
        pem_rfc7468::decode_vec(public_key_maybe_pem.as_bytes())
          .map_err(anyhow::Error::msg)
          .context("Failed to get der from pem")?;
      Ok(public_key_der)
    } else {
      BASE64
        .decode(public_key_maybe_pem.as_bytes())
        .context("Public key is not base64")
    }
  }

  pub fn der_to_raw_bytes(
    spki_der: &[u8],
  ) -> anyhow::Result<[u8; 32]> {
    let Ok(spki) = spki::SubjectPublicKeyInfo::<
      der::AnyRef<'_>,
      BitStringRef<'_>,
    >::try_from(spki_der) else {
      return Err(anyhow!("Public key is not der"));
    };

    if spki.algorithm.oid != super::OID_X25519 {
      return Err(anyhow!("Public key is not X25519"));
    }

    let bs = spki.subject_public_key;

    // Check byte aligned
    if bs.unused_bits() != 0 {
      return Err(anyhow!("Public key spki der is not byte-aligned"));
    }

    let raw = bs.as_bytes().context("Bitstring has no bytes")?;
    if raw.len() != 32 {
      return Err(anyhow!(
        "Raw public key length should be 32, got {}",
        raw.len()
      ));
    }

    let mut res = [0u8; 32];
    res.copy_from_slice(raw);
    Ok(res)
  }

  pub fn from_private_key_using_dh(
    pki_kind: PkiKind,
    maybe_pkcs8_private_key: &str,
  ) -> anyhow::Result<Self> {
    let params: NoiseParams = pki_kind.noise_params().parse()?;
    let mut dh = DefaultResolver.resolve_dh(&params.dh).context(
      "No DH implementation available with these noise params",
    )?;
    let private_key = crate::key::Pkcs8PrivateKey::maybe_raw_bytes(
      maybe_pkcs8_private_key,
    )?;
    dh.set(&private_key);
    let raw_public_key = dh.pubkey().to_vec();
    Self::from_raw_bytes(&raw_public_key)
  }
}
