use byteorder::{BigEndian, ReadBytesExt as _, WriteBytesExt as _};
use std::ffi::OsStr;
use std::io::{Read as _, Write};
use std::os::unix::fs::PermissionsExt as _;
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
use uuid::Uuid;

#[repr(u8)]
pub enum TlvTag {
    CryptoPrivateK = 0x01,
    EncryptedDataC1 = 0x02,
    EncryptedDataCt = 0x03,
    EncryptedDataIv = 0x04,
    End = 0xFF,
}

impl TlvTag {
    pub fn from_u8(value: u8) -> Option<Self> {
        let enum_value = if value == Self::CryptoPrivateK as u8 {
            Self::CryptoPrivateK
        } else if value == Self::EncryptedDataC1 as u8 {
            Self::EncryptedDataC1
        } else if value == Self::EncryptedDataCt as u8 {
            Self::EncryptedDataCt
        } else if value == Self::EncryptedDataIv as u8 {
            Self::EncryptedDataIv
        } else if value == Self::End as u8 {
            Self::End
        } else {
            return None;
        };
        Some(enum_value)
    }
}

#[derive(Clone)]
pub struct PatchData {
    pub private_key: crypto::PrivateKey,
    pub encrypted_code: crypto::Ciphertext,
}

#[derive(Default)]
pub struct BankStorage {
    root: PathBuf,
}

impl BankStorage {
    pub fn new(root: impl AsRef<Path>) -> anyhow::Result<Self> {
        let root = root.as_ref().to_path_buf();
        std::fs::create_dir_all(&root)?;
        Ok(Self { root })
    }

    fn patch_path(&self, id: Uuid) -> PathBuf {
        self.root.join(format!("{}.pb", id))
    }

    pub fn put_patch(&self, id: Uuid, patch: &PatchData) -> anyhow::Result<()> {
        self.write_patch_file(&self.patch_path(id), patch)
    }

    pub fn patches(
        &self,
    ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<PatchData>> + use<'_>> {
        let iter = std::fs::read_dir(self.root.clone())?.map(|maybe_dir| {
            maybe_dir
                .map_err(anyhow::Error::from)
                .and_then(|dir| Ok(Uuid::try_parse_ascii(dir.file_name().as_encoded_bytes())?))
                .and_then(|id| self.get_patch(id))
        });
        Ok(iter)
    }

    pub fn get_patch(&self, id: Uuid) -> anyhow::Result<PatchData> {
        self.read_patch_file(&self.root.join(format!("{}.pb", id)))
    }

    pub fn write_patch_file(&self, path: &Path, patch: &PatchData) -> anyhow::Result<()> {
        {
            let mut file = std::fs::File::create(path)?;
            write_tlv(
                &mut file,
                TlvTag::CryptoPrivateK,
                &patch.private_key.k().to_be_bytes(),
            )?;
            write_tlv(
                &mut file,
                TlvTag::EncryptedDataC1,
                &patch.encrypted_code.c1.value().to_be_bytes(),
            )?;
            write_tlv(&mut file, TlvTag::EncryptedDataCt, &patch.encrypted_code.ct)?;
            write_tlv(&mut file, TlvTag::EncryptedDataIv, &patch.encrypted_code.iv)?;
            file.write_all(&[TlvTag::End as u8])?;
        }
        let permissions = std::fs::Permissions::from_mode(0o600);
        std::fs::set_permissions(path, permissions)?;
        Ok(())
    }

    pub fn read_patch_file(&self, path: &Path) -> anyhow::Result<PatchData> {
        let mut file = std::fs::File::open(path)?;

        let mut private_key = None;
        let mut c1 = None;
        let mut ct = None;
        let mut iv = None;
        while let Ok(tag) = file.read_u8() {
            if tag == TlvTag::End as u8 {
                break;
            }
            let len = file.read_u16::<BigEndian>()? as usize;
            let mut v = vec![0u8; len];
            file.read_exact(&mut v)?;
            match TlvTag::from_u8(tag).ok_or(anyhow::anyhow!("Unknown tag"))? {
                TlvTag::CryptoPrivateK => {
                    private_key = Some(crypto::PrivateKey::from_k(
                        &crypto::DEFAULT_CRYPTO_PARAMS,
                        crypto::U512::from_be_slice(&v),
                    )?)
                }
                TlvTag::EncryptedDataC1 => {
                    c1 = Some(
                        crypto::DEFAULT_CRYPTO_PARAMS
                            .ring()
                            .reduce(crypto::U512::from_be_slice(&v).resize()),
                    )
                }
                TlvTag::EncryptedDataCt => ct = Some(v),
                TlvTag::EncryptedDataIv => {
                    iv = Some(
                        v.try_into()
                            .map_err(|_| anyhow::anyhow!("IV len mismatch"))?,
                    )
                }
                TlvTag::End => anyhow::bail!("Should be unreachable"),
            }
        }
        let Some(private_key) = private_key else {
            anyhow::bail!("private key missing");
        };
        let Some(c1) = c1 else {
            anyhow::bail!("c1 missing");
        };
        let Some(ct) = ct else {
            anyhow::bail!("ct missing");
        };
        let Some(iv) = iv else {
            anyhow::bail!("iv missing");
        };
        Ok(PatchData {
            private_key,
            encrypted_code: crypto::Ciphertext { c1, ct, iv },
        })
    }

    pub fn clean_old_patches(&self, max_age: Duration) -> anyhow::Result<usize> {
        let cutoff = SystemTime::now()
            .checked_sub(max_age)
            .ok_or_else(|| anyhow::anyhow!("max_age too large"))?;

        let mut removed = 0usize;
        for entry in std::fs::read_dir(&self.root)? {
            let entry = entry?;
            let path = entry.path();

            if path.extension().and_then(OsStr::to_str) != Some("pb") {
                continue;
            }
            let Some(stem) = path.file_stem().and_then(|s| s.to_str()) else {
                continue;
            };
            if Uuid::parse_str(stem).is_err() {
                continue;
            }

            let metadata = entry.metadata()?;
            let modified = metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH);
            if modified <= cutoff {
                std::fs::remove_file(&path)?;
                removed += 1;
            }
        }
        Ok(removed)
    }
}

fn write_tlv(mut out: impl Write, tag: TlvTag, value: &[u8]) -> anyhow::Result<()> {
    out.write_all(&[tag as u8])?;
    out.write_u16::<BigEndian>(value.len() as u16)?;
    out.write_all(value)?;
    Ok(())
}
