use pyo3::prelude::*;
use pyo3::types::PyBytes;
use secp256k1::{rand, SecretKey, SECP256K1};

use crate::message_from_digest;
use crate::public_key::PublicKey;

#[pyclass]
pub struct PrivateKey {
    secret: [u8; 32],
}

#[pymethods]
impl PrivateKey {
    #[new]
    #[pyo3(signature = (secret = None))]
    fn new(secret: Option<&[u8]>) -> PyResult<Self> {
        let secret = match secret {
            Some(bytes) => {
                // Validate it's a valid secret key
                let sk = SecretKey::from_slice(bytes)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
                sk.secret_bytes()
            }
            None => {
                let sk = SecretKey::new(&mut rand::thread_rng());
                sk.secret_bytes()
            }
        };
        Ok(Self { secret })
    }

    fn sign_recoverable<'py>(
        &self,
        py: Python<'py>,
        message: &[u8],
    ) -> PyResult<Bound<'py, PyBytes>> {
        let sk = SecretKey::from_slice(&self.secret)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;

        let msg = message_from_digest(message)?;

        let sig = SECP256K1.sign_ecdsa_recoverable(&msg, &sk);
        let (recid, compact) = sig.serialize_compact();

        let mut out = [0u8; 65];
        out[..64].copy_from_slice(&compact);
        out[64] = i32::from(recid) as u8;

        Ok(PyBytes::new(py, &out))
    }

    #[getter]
    fn public_key(&self) -> PyResult<PublicKey> {
        let sk = SecretKey::from_slice(&self.secret)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
        let pk = secp256k1::PublicKey::from_secret_key(SECP256K1, &sk);
        Ok(PublicKey::from_inner(pk))
    }
}
