use pyo3::prelude::*;
use pyo3::types::PyBytes;
use secp256k1::{ecdsa::RecoverableSignature, SECP256K1};

use crate::message_from_digest;

#[pyclass]
pub struct PublicKey {
    inner: secp256k1::PublicKey,
}

impl PublicKey {
    pub fn from_inner(inner: secp256k1::PublicKey) -> Self {
        Self { inner }
    }
}

#[pymethods]
impl PublicKey {
    #[new]
    fn new(data: &[u8]) -> PyResult<Self> {
        let inner = secp256k1::PublicKey::from_slice(data)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        Ok(Self { inner })
    }

    #[staticmethod]
    fn from_signature_and_message(
        signature: &[u8],
        message: &[u8],
    ) -> PyResult<Self> {
        if signature.len() != 65 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "signature must be 65 bytes",
            ));
        }

        let recid = secp256k1::ecdsa::RecoveryId::try_from(signature[64] as i32)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

        let sig = RecoverableSignature::from_compact(&signature[..64], recid)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

        let msg = message_from_digest(message)?;

        let inner = SECP256K1
            .recover_ecdsa(&msg, &sig)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

        Ok(Self { inner })
    }

    #[staticmethod]
    fn from_secret(secret: &[u8]) -> PyResult<Self> {
        let sk = secp256k1::SecretKey::from_slice(secret)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        let inner = secp256k1::PublicKey::from_secret_key(SECP256K1, &sk);
        Ok(Self { inner })
    }

    #[pyo3(signature = (compressed = true))]
    fn format<'py>(&self, py: Python<'py>, compressed: bool) -> Bound<'py, PyBytes> {
        if compressed {
            PyBytes::new(py, &self.inner.serialize())
        } else {
            PyBytes::new(py, &self.inner.serialize_uncompressed())
        }
    }
}
