use crate::{AuditEnvelope, Invariant, InvariantResult, State, TransitionParams};
use serde_json::json;
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use time::OffsetDateTime;
use uuid::Uuid;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum DteError {
    #[error("Invariant violation: {0:?}")]
    InvariantViolation(Vec<String>),
}

pub type TransitionFn =
    fn(&State, &TransitionParams) -> HashMap<String, Value>;

fn canonical_json(value: &Value) -> String {
    let v = value.clone();
    // serde_json already produces deterministic ordering for maps when using Value
    serde_json::to_string(&v).expect("canonical_json serialization failed")
}

fn sha256_hex(s: &str) -> String {
    let mut hasher = Sha256::new();
    hasher.update(s.as_bytes());
    format!("{:x}", hasher.finalize())
}

pub struct DeterministicTransitionEngine {
    invariants: Vec<Box<dyn Invariant>>,
    signer_key_id: String,
    pub prev_envelope_hash: Option<String>,
}

impl DeterministicTransitionEngine {
    pub fn new(
        invariants: Vec<Box<dyn Invariant>>,
        signer_key_id: String,
    ) -> Self {
        Self {
            invariants,
            signer_key_id,
            prev_envelope_hash: None,
        }
    }

    fn compute_state_hash(&self, state: &State) -> String {
        let value = serde_json::to_value(&state.data).unwrap();
        sha256_hex(&canonical_json(&value))
    }

    fn compute_transcript_hash(
        &self,
        before: &State,
        after: &State,
        params: &TransitionParams,
        result: &HashMap<String, Value>,
    ) -> String {
        let transcript = json!({
            "before": before.data,
            "after": after.data,
            "params": {
                "action": params.action,
                "args": params.args,
                "identity_subject": params.identity_subject,
                "identity_roles": params.identity_roles,
                "giu_cost": params.giu_cost,
            },
            "result": result,
        });
        sha256_hex(&canonical_json(&transcript))
    }

    fn sign(&self, transcript_hash: &str) -> String {
        // K-1 placeholder signature; replace with real crypto later
        sha256_hex(&format!("{}:{}", transcript_hash, self.signer_key_id))
    }

    fn check_invariants(
        &self,
        before: &State,
        after: &State,
        params: &TransitionParams,
    ) -> InvariantResult {
        let mut failed = Vec::new();
        for inv in &self.invariants {
            if !inv.check(before, after, params) {
                failed.push(inv.name().to_string());
            }
        }
        if failed.is_empty() {
            InvariantResult::success()
        } else {
            InvariantResult::failure(failed)
        }
    }

    pub fn apply(
        &mut self,
        current_state: &State,
        params: &TransitionParams,
        transition_fn: TransitionFn,
        request_id: Option<String>,
    ) -> Result<(State, AuditEnvelope), DteError> {
        // 1. pure deterministic result
        let result = transition_fn(current_state, params);

        // 2. derive next_state (simple overlay)
        let mut next_data = current_state.data.clone();
        if let Some(updates) = result.get("state_updates") {
            if let Some(map) = updates.as_object() {
                for (k, v) in map {
                    next_data.insert(k.clone(), v.clone());
                }
            }
        }
        let next_state = State::new(next_data);

        // 3. hashes
        let state_hash_before = self.compute_state_hash(current_state);
        let state_hash_after = self.compute_state_hash(&next_state);
        let transcript_hash = self.compute_transcript_hash(
            current_state,
            &next_state,
            params,
            &result,
        );

        // 4. invariants
        let inv_result = self.check_invariants(current_state, &next_state, params);
        if !inv_result.ok {
            return Err(DteError::InvariantViolation(inv_result.failed));
        }

        // 5. signature + envelope core
        let signature = self.sign(&transcript_hash);
        let id = Uuid::new_v4().to_string();
        let req_id = request_id.unwrap_or_else(|| Uuid::new_v4().to_string());
        let created_at = OffsetDateTime::now_utc()
            .format(&time::format_description::well_known::Rfc3339)
            .unwrap();

        let signer = {
            let mut m = HashMap::new();
            m.insert("key_id".to_string(), self.signer_key_id.clone());
            m
        };

        let payload = params.args.clone();
        let result_clone = result.clone();

        let mut core = AuditEnvelope {
            id,
            request_id: req_id,
            action: params.action.clone(),
            identity_subject: params.identity_subject.clone(),
            identity_roles: params.identity_roles.clone(),
            payload,
            result: result_clone,
            giu_cost: params.giu_cost,
            state_hash_before: Some(state_hash_before),
            state_hash_after: Some(state_hash_after),
            created_at,
            transcript_hash,
            signature,
            signer,
            prev_envelope_hash: self.prev_envelope_hash.clone(),
            record_hash: String::new(), // fill after hash
        };

        let core_value = serde_json::to_value(&core).unwrap();
        let record_hash = sha256_hex(&canonical_json(&core_value));
        core.record_hash = record_hash.clone();

        // 6. update chain head
        self.prev_envelope_hash = Some(record_hash);

        Ok((next_state, core))
    }
}
