//! UTXO (Unspent Transaction Output) management for DAG-based consensus.
//!
//! In a DAG, multiple blocks can reference the same parent and potentially
//! contain conflicting transactions. The UTXO model must:
//!
//! 1. Track UTXOs per virtual block state (not global)
//! 2. Support efficient diff-based updates
//! 3. Handle transaction conflicts via block ordering

use hashbrown::{HashMap, HashSet};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use synor_types::{
    transaction::{Outpoint, ScriptPubKey, Transaction, TxOutput},
    Address, Amount, Hash256,
};
use thiserror::Error;

/// A single UTXO entry.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct UtxoEntry {
    /// The output data.
    pub output: TxOutput,
    /// Block in which this UTXO was created.
    pub block_daa_score: u64,
    /// Whether this is a coinbase output.
    pub is_coinbase: bool,
}

impl UtxoEntry {
    /// Creates a new UTXO entry.
    pub fn new(output: TxOutput, block_daa_score: u64, is_coinbase: bool) -> Self {
        UtxoEntry {
            output,
            block_daa_score,
            is_coinbase,
        }
    }

    /// Returns the amount.
    pub fn amount(&self) -> Amount {
        self.output.amount
    }

    /// Returns the script pubkey.
    pub fn script_pubkey(&self) -> &ScriptPubKey {
        &self.output.script_pubkey
    }

    /// Checks if this UTXO is mature enough to spend.
    pub fn is_mature(&self, current_daa_score: u64) -> bool {
        if !self.is_coinbase {
            return true;
        }
        // Coinbase must wait for maturity period
        current_daa_score >= self.block_daa_score + crate::COINBASE_MATURITY
    }
}

/// Diff representing changes to UTXO set.
#[derive(Clone, Debug, Default)]
pub struct UtxoDiff {
    /// UTXOs to add.
    pub to_add: HashMap<Outpoint, UtxoEntry>,
    /// UTXOs to remove.
    pub to_remove: HashSet<Outpoint>,
}

impl UtxoDiff {
    /// Creates a new empty diff.
    pub fn new() -> Self {
        UtxoDiff {
            to_add: HashMap::new(),
            to_remove: HashSet::new(),
        }
    }

    /// Adds a UTXO to the diff.
    pub fn add(&mut self, outpoint: Outpoint, entry: UtxoEntry) {
        // If it was marked for removal, cancel that
        self.to_remove.remove(&outpoint);
        self.to_add.insert(outpoint, entry);
    }

    /// Removes a UTXO in the diff.
    pub fn remove(&mut self, outpoint: Outpoint) {
        // If it was marked for addition, cancel that (removal cancels the add)
        if self.to_add.remove(&outpoint).is_none() {
            // Only mark for removal if it wasn't in to_add
            // (i.e., it's a pre-existing UTXO being removed)
            self.to_remove.insert(outpoint);
        }
    }

    /// Merges another diff into this one.
    pub fn merge(&mut self, other: UtxoDiff) {
        // Apply other's removals
        for outpoint in other.to_remove {
            self.remove(outpoint);
        }
        // Apply other's additions
        for (outpoint, entry) in other.to_add {
            self.add(outpoint, entry);
        }
    }

    /// Returns true if the diff is empty.
    pub fn is_empty(&self) -> bool {
        self.to_add.is_empty() && self.to_remove.is_empty()
    }

    /// Returns the reverse of this diff.
    pub fn reverse(&self) -> UtxoDiff {
        let mut reversed = UtxoDiff::new();
        // Additions become removals
        for outpoint in self.to_add.keys() {
            reversed.to_remove.insert(*outpoint);
        }
        // Removals... we can't fully reverse without the original entries
        // This is a limitation - in practice we'd need to store the removed entries
        reversed
    }
}

/// The main UTXO set.
pub struct UtxoSet {
    /// Current UTXOs.
    utxos: RwLock<HashMap<Outpoint, UtxoEntry>>,
    /// Total value in the set.
    total_value: RwLock<Amount>,
    /// Number of UTXOs.
    count: RwLock<usize>,
}

impl UtxoSet {
    /// Creates a new empty UTXO set.
    pub fn new() -> Self {
        UtxoSet {
            utxos: RwLock::new(HashMap::new()),
            total_value: RwLock::new(Amount::ZERO),
            count: RwLock::new(0),
        }
    }

    /// Creates a UTXO set from existing entries.
    pub fn from_entries(entries: HashMap<Outpoint, UtxoEntry>) -> Self {
        let total: Amount = entries
            .values()
            .map(|e| e.amount())
            .fold(Amount::ZERO, |a, b| a.saturating_add(b));
        let count = entries.len();
        UtxoSet {
            utxos: RwLock::new(entries),
            total_value: RwLock::new(total),
            count: RwLock::new(count),
        }
    }

    /// Gets a UTXO by outpoint.
    pub fn get(&self, outpoint: &Outpoint) -> Option<UtxoEntry> {
        self.utxos.read().get(outpoint).cloned()
    }

    /// Checks if a UTXO exists.
    pub fn contains(&self, outpoint: &Outpoint) -> bool {
        self.utxos.read().contains_key(outpoint)
    }

    /// Returns the number of UTXOs.
    pub fn len(&self) -> usize {
        *self.count.read()
    }

    /// Returns true if empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the total value in the set.
    pub fn total_value(&self) -> Amount {
        *self.total_value.read()
    }

    /// Adds a UTXO to the set.
    pub fn add(&self, outpoint: Outpoint, entry: UtxoEntry) -> Result<(), UtxoError> {
        let mut utxos = self.utxos.write();
        let mut total = self.total_value.write();
        let mut count = self.count.write();

        if utxos.contains_key(&outpoint) {
            return Err(UtxoError::UtxoAlreadyExists(outpoint));
        }

        *total = total.saturating_add(entry.amount());
        *count += 1;
        utxos.insert(outpoint, entry);

        Ok(())
    }

    /// Removes a UTXO from the set.
    pub fn remove(&self, outpoint: &Outpoint) -> Result<UtxoEntry, UtxoError> {
        let mut utxos = self.utxos.write();
        let mut total = self.total_value.write();
        let mut count = self.count.write();

        let entry = utxos
            .remove(outpoint)
            .ok_or(UtxoError::UtxoNotFound(*outpoint))?;

        *total = total.saturating_sub(entry.amount());
        *count -= 1;

        Ok(entry)
    }

    /// Applies a diff to the UTXO set.
    pub fn apply_diff(&self, diff: &UtxoDiff) -> Result<(), UtxoError> {
        // Apply removals first
        for outpoint in &diff.to_remove {
            if !self.contains(outpoint) {
                return Err(UtxoError::UtxoNotFound(*outpoint));
            }
        }

        // Check additions don't conflict
        for outpoint in diff.to_add.keys() {
            if self.contains(outpoint) {
                return Err(UtxoError::UtxoAlreadyExists(*outpoint));
            }
        }

        // Now apply (we've validated above)
        for outpoint in &diff.to_remove {
            self.remove(outpoint)?;
        }

        for (outpoint, entry) in &diff.to_add {
            self.add(*outpoint, entry.clone())?;
        }

        Ok(())
    }

    /// Creates a diff from applying a transaction.
    pub fn create_transaction_diff(
        &self,
        tx: &Transaction,
        block_daa_score: u64,
    ) -> Result<UtxoDiff, UtxoError> {
        let mut diff = UtxoDiff::new();
        let txid = tx.txid();

        // Process inputs (remove UTXOs)
        if !tx.is_coinbase() {
            for input in &tx.inputs {
                let outpoint = input.previous_output;

                // Verify the UTXO exists
                let entry = self
                    .get(&outpoint)
                    .ok_or(UtxoError::UtxoNotFound(outpoint))?;

                // Check maturity for coinbase inputs
                if !entry.is_mature(block_daa_score) {
                    return Err(UtxoError::ImmatureCoinbase(outpoint));
                }

                diff.remove(outpoint);
            }
        }

        // Process outputs (add UTXOs)
        for (index, output) in tx.outputs.iter().enumerate() {
            let outpoint = Outpoint::new(txid, index as u32);
            let entry = UtxoEntry::new(output.clone(), block_daa_score, tx.is_coinbase());
            diff.add(outpoint, entry);
        }

        Ok(diff)
    }

    /// Gets all UTXOs for a given address.
    pub fn get_by_address(
        &self,
        address: &Address,
        network: synor_types::Network,
    ) -> Vec<(Outpoint, UtxoEntry)> {
        let utxos = self.utxos.read();
        let mut result = Vec::new();

        for (outpoint, entry) in utxos.iter() {
            if let Some(utxo_address) = entry.script_pubkey().extract_address(network) {
                if &utxo_address == address {
                    result.push((*outpoint, entry.clone()));
                }
            }
        }

        result
    }

    /// Gets the balance for an address.
    pub fn get_balance(&self, address: &Address, network: synor_types::Network) -> Amount {
        self.get_by_address(address, network)
            .iter()
            .map(|(_, entry)| entry.amount())
            .fold(Amount::ZERO, |a, b| a.saturating_add(b))
    }

    /// Selects UTXOs for spending a target amount.
    pub fn select_utxos(
        &self,
        address: &Address,
        target: Amount,
        network: synor_types::Network,
    ) -> Result<Vec<(Outpoint, UtxoEntry)>, UtxoError> {
        let available = self.get_by_address(address, network);

        // Sort by amount descending for better selection
        let mut sorted = available;
        sorted.sort_by(|a, b| b.1.amount().cmp(&a.1.amount()));

        let mut selected = Vec::new();
        let mut total = Amount::ZERO;

        for (outpoint, entry) in sorted {
            if total >= target {
                break;
            }
            total = total.saturating_add(entry.amount());
            selected.push((outpoint, entry));
        }

        if total < target {
            return Err(UtxoError::InsufficientFunds {
                required: target,
                available: total,
            });
        }

        Ok(selected)
    }

    /// Creates a snapshot of the current UTXO set.
    pub fn snapshot(&self) -> HashMap<Outpoint, UtxoEntry> {
        self.utxos.read().clone()
    }

    /// Clears all UTXOs.
    pub fn clear(&self) {
        let mut utxos = self.utxos.write();
        let mut total = self.total_value.write();
        let mut count = self.count.write();

        utxos.clear();
        *total = Amount::ZERO;
        *count = 0;
    }
}

impl Default for UtxoSet {
    fn default() -> Self {
        Self::new()
    }
}

impl std::fmt::Debug for UtxoSet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UtxoSet")
            .field("count", &self.len())
            .field("total_value", &self.total_value())
            .finish()
    }
}

/// Virtual UTXO set for DAG-based tracking.
///
/// Tracks UTXO state at each block position in the DAG.
pub struct VirtualUtxoState {
    /// Base UTXO set.
    base: UtxoSet,
    /// Pending diffs per block.
    pending_diffs: RwLock<HashMap<Hash256, UtxoDiff>>,
    /// Current virtual block (tip of selected chain).
    virtual_block: RwLock<Hash256>,
}

impl VirtualUtxoState {
    /// Creates a new virtual UTXO state.
    pub fn new(genesis_hash: Hash256) -> Self {
        VirtualUtxoState {
            base: UtxoSet::new(),
            pending_diffs: RwLock::new(HashMap::new()),
            virtual_block: RwLock::new(genesis_hash),
        }
    }

    /// Returns the current virtual block.
    pub fn virtual_block(&self) -> Hash256 {
        *self.virtual_block.read()
    }

    /// Returns the base UTXO set.
    pub fn base(&self) -> &UtxoSet {
        &self.base
    }

    /// Adds a block's diff.
    pub fn add_block_diff(&self, block_hash: Hash256, diff: UtxoDiff) {
        self.pending_diffs.write().insert(block_hash, diff);
    }

    /// Gets a block's diff.
    pub fn get_block_diff(&self, block_hash: &Hash256) -> Option<UtxoDiff> {
        self.pending_diffs.read().get(block_hash).cloned()
    }

    /// Updates the virtual block and applies necessary diffs.
    pub fn update_virtual(
        &self,
        new_virtual: Hash256,
        chain_diffs: Vec<UtxoDiff>,
    ) -> Result<(), UtxoError> {
        // Apply all diffs in order
        for diff in chain_diffs {
            self.base.apply_diff(&diff)?;
        }

        *self.virtual_block.write() = new_virtual;
        Ok(())
    }

    /// Checks if a UTXO is available at the current virtual state.
    pub fn is_available(&self, outpoint: &Outpoint) -> bool {
        self.base.contains(outpoint)
    }
}

impl std::fmt::Debug for VirtualUtxoState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("VirtualUtxoState")
            .field("virtual_block", &self.virtual_block())
            .field("pending_diffs", &self.pending_diffs.read().len())
            .finish()
    }
}

/// Errors related to UTXOs.
#[derive(Debug, Error)]
pub enum UtxoError {
    #[error("UTXO not found: {0}")]
    UtxoNotFound(Outpoint),

    #[error("UTXO already exists: {0}")]
    UtxoAlreadyExists(Outpoint),

    #[error("Immature coinbase output: {0}")]
    ImmatureCoinbase(Outpoint),

    #[error("Insufficient funds: required {required}, available {available}")]
    InsufficientFunds { required: Amount, available: Amount },

    #[error("Double spend detected: {0}")]
    DoubleSpend(Outpoint),

    #[error("Invalid transaction: {0}")]
    InvalidTransaction(String),
}

#[cfg(test)]
mod tests {
    use super::*;
    use synor_types::transaction::ScriptPubKey;

    fn make_outpoint(n: u8) -> Outpoint {
        let mut bytes = [0u8; 32];
        bytes[0] = n;
        Outpoint::new(Hash256::from_bytes(bytes), 0)
    }

    fn make_entry(amount: u64) -> UtxoEntry {
        UtxoEntry::new(
            TxOutput::new(Amount::from_sompi(amount), ScriptPubKey::p2pkh(&[0u8; 32])),
            100,
            false,
        )
    }

    fn make_coinbase_entry(amount: u64, daa_score: u64) -> UtxoEntry {
        UtxoEntry::new(
            TxOutput::new(Amount::from_sompi(amount), ScriptPubKey::p2pkh(&[0u8; 32])),
            daa_score,
            true,
        )
    }

    #[test]
    fn test_utxo_set_basic() {
        let set = UtxoSet::new();
        assert!(set.is_empty());
        assert_eq!(set.total_value(), Amount::ZERO);

        let outpoint = make_outpoint(1);
        let entry = make_entry(1000);

        set.add(outpoint, entry.clone()).unwrap();
        assert_eq!(set.len(), 1);
        assert_eq!(set.total_value(), Amount::from_sompi(1000));

        let retrieved = set.get(&outpoint).unwrap();
        assert_eq!(retrieved.amount(), entry.amount());
    }

    #[test]
    fn test_utxo_set_remove() {
        let set = UtxoSet::new();

        let outpoint = make_outpoint(1);
        let entry = make_entry(1000);

        set.add(outpoint, entry).unwrap();
        assert_eq!(set.len(), 1);

        set.remove(&outpoint).unwrap();
        assert!(set.is_empty());
        assert_eq!(set.total_value(), Amount::ZERO);
    }

    #[test]
    fn test_utxo_diff() {
        let set = UtxoSet::new();

        // Add initial UTXO
        let outpoint1 = make_outpoint(1);
        set.add(outpoint1, make_entry(1000)).unwrap();

        // Create diff: remove 1, add 2
        let mut diff = UtxoDiff::new();
        diff.remove(outpoint1);

        let outpoint2 = make_outpoint(2);
        diff.add(outpoint2, make_entry(2000));

        // Apply diff
        set.apply_diff(&diff).unwrap();

        assert!(!set.contains(&outpoint1));
        assert!(set.contains(&outpoint2));
        assert_eq!(set.total_value(), Amount::from_sompi(2000));
    }

    #[test]
    fn test_coinbase_maturity() {
        let entry = make_coinbase_entry(5000, 100);

        // Not mature at same score
        assert!(!entry.is_mature(100));

        // Not mature just before threshold
        assert!(!entry.is_mature(199));

        // Mature at exactly maturity threshold
        assert!(entry.is_mature(200));

        // Mature after threshold
        assert!(entry.is_mature(300));
    }

    #[test]
    fn test_non_coinbase_always_mature() {
        let entry = make_entry(1000);
        assert!(entry.is_mature(0));
        assert!(entry.is_mature(100));
    }

    #[test]
    fn test_duplicate_add_error() {
        let set = UtxoSet::new();
        let outpoint = make_outpoint(1);

        set.add(outpoint, make_entry(1000)).unwrap();

        let result = set.add(outpoint, make_entry(2000));
        assert!(matches!(result, Err(UtxoError::UtxoAlreadyExists(_))));
    }

    #[test]
    fn test_remove_nonexistent_error() {
        let set = UtxoSet::new();
        let outpoint = make_outpoint(99);

        let result = set.remove(&outpoint);
        assert!(matches!(result, Err(UtxoError::UtxoNotFound(_))));
    }

    #[test]
    fn test_diff_merge() {
        let mut diff1 = UtxoDiff::new();
        diff1.add(make_outpoint(1), make_entry(100));
        diff1.add(make_outpoint(2), make_entry(200));

        let mut diff2 = UtxoDiff::new();
        diff2.remove(make_outpoint(1)); // Remove what diff1 added
        diff2.add(make_outpoint(3), make_entry(300));

        diff1.merge(diff2);

        // Outpoint 1 should be cancelled out
        assert!(!diff1.to_add.contains_key(&make_outpoint(1)));
        assert!(!diff1.to_remove.contains(&make_outpoint(1)));

        // Outpoint 2 and 3 should be added
        assert!(diff1.to_add.contains_key(&make_outpoint(2)));
        assert!(diff1.to_add.contains_key(&make_outpoint(3)));
    }
}
