#![allow(unsafe_op_in_unsafe_fn)]

use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};
use pyo3::exceptions::PyValueError;
use sev::firmware::guest::*;
use sev::parser::ByteParser;


fn to_py_value_err(msg: impl Into<String>) -> PyErr {
    /* 
    Convert Rust errors into Python ValueError.
    */
    PyValueError::new_err(msg.into())
}


#[pyfunction]
fn get_report(py: Python<'_>, challenge: &[u8], msg_ver: Option<u32>, vmpl: Option<u32>) -> PyResult<Py<PyBytes>> {
    /*
    Fetch an AMD SEV-SNP attestation report from the local platform.

    The report call will:
    - Enforce challenge is exactly 64 bytes (mapped to the SNP report_data field)
    - Invoke the SEV-SNP firmware interface to obtain a report
    - Return the raw attestation report bytes to Python

    Args:
    - challenge: Nonce buffer inserted into report_data (must be exactly 64 bytes).
    - msg_ver: Optional message/report version selector (defaults to 1).
    - vmpl: Optional VM privilege level selector (defaults to 1).

    Returns:
    - Raw SEV-SNP attestation report bytes
    */
    let msg_ver = msg_ver.unwrap_or(1);
    let vmpl = vmpl.unwrap_or(1);

    // sev 7.1.0 expects a fixed 64-byte challenge
    if challenge.len() != 64 {
        return Err(to_py_value_err(format!(
            "challenge must be exactly 64 bytes, got {}",
            challenge.len()
        )));
    }

    let mut challenge_arr = [0u8; 64];
    challenge_arr.copy_from_slice(challenge);

    let mut fw = Firmware::open()
        .map_err(|e| PyValueError::new_err(format!("Firmware::open failed: {e:?}")))?;

    let report: Vec<u8> = fw
        .get_report(Some(msg_ver), Some(challenge_arr), Some(vmpl))
        .map_err(|e| PyValueError::new_err(format!("get_report failed: {e:?}")))?;

    Ok(PyBytes::new_bound(py, &report).unbind())
}


#[pyfunction]
fn parse_report(py: Python<'_>, report_bytes: &[u8]) -> PyResult<PyObject> {
    /*
    Parse an AMD SEV-SNP attestation report into structured fields for Python.

    The parser will:
    - Decode report_bytes using the sev library authoritative ABI parser
    - Extract report_data, chip_id, reported_tcb, and CPUID fields
    - Extract the signed prefix bytes and ECDSA (r, s) signature components
    - Return a Python dict compatible with ParsedAttestationReport construction

    Args:
    - py: Python interpreter token for allocating Python objects.
    - report_bytes: Raw SEV-SNP attestation report bytes.

    Returns:
    - Python dict containing parsed report fields and signature material
    */
    let report = AttestationReport::from_bytes(report_bytes)
        .map_err(|e| to_py_value_err(format!("AttestationReport::from_bytes failed: {e:?}")))?;

    // Match the Rust verifier exactly: take signed_bytes from report.to_bytes()
    let serialized = report.to_bytes().map_err(|e| to_py_value_err(format!("report.to_bytes failed: {e:?}")))?;

    // The signature covers bytes 0..0x29F inclusive => slice [0..0x2A0)
    if serialized.len() < 0x2A0 {
        return Err(to_py_value_err(format!(
            "serialized report too short for signed prefix: got {}, need at least {}",
            serialized.len(),
            0x2A0
        )));
    }
    let signed_bytes = &serialized[0..0x2A0];

    // Extract ECDSA signature components from the report
    let r = report.signature.r();
    let s = report.signature.s();

    // Build Python dict
    let out = PyDict::new_bound(py);
    out.set_item("chip_id", PyBytes::new_bound(py, &report.chip_id))?;
    out.set_item("report_data", PyBytes::new_bound(py, &report.report_data))?;
    out.set_item("signed_bytes", PyBytes::new_bound(py, signed_bytes))?;
    out.set_item("sig_r", PyBytes::new_bound(py, r.as_slice()))?;
    out.set_item("sig_s", PyBytes::new_bound(py, s.as_slice()))?;

    // reported_tcb fields
    let tcb = PyDict::new_bound(py);
    tcb.set_item("bootloader", report.reported_tcb.bootloader)?;
    tcb.set_item("tee", report.reported_tcb.tee)?;
    tcb.set_item("snp", report.reported_tcb.snp)?;
    tcb.set_item("microcode", report.reported_tcb.microcode)?;
    out.set_item("reported_tcb", tcb)?;

    // cpuid fields
    let cpu_info = PyDict::new_bound(py);
    cpu_info.set_item("fam_id", report.cpuid_fam_id)?;
    cpu_info.set_item("mod_id", report.cpuid_mod_id)?;
    cpu_info.set_item("step", report.cpuid_step)?;
    out.set_item("cpuid", cpu_info)?;

    Ok(out.into())
}


#[pymodule]
fn tee_amd_sev_snp_rust_py_wrap(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    /*
    Register the Rust-backed SEV-SNP helpers as a Python module.

    The module will:
    - Expose get_report() for generating a raw SEV-SNP attestation report
    - Expose parse_report() for decoding report bytes into structured fields

    Args:
    - _py: Python interpreter token for module initialization.
    - m: Python module handle used to register functions.

    Returns:
    - Ok(()) if module initialization succeeds
    */
    m.add_function(wrap_pyfunction!(get_report, m)?)?;
    m.add_function(wrap_pyfunction!(parse_report, m)?)?;
    Ok(())
}