"""
BloomOS Kernel - IT Token Ledger
================================

Unified resource accounting for all syscalls.

IT (Information Token) is the universal compute currency:
- CPU time costs IT
- I/O bandwidth costs IT
- RAM allocation costs IT

The ledger is a tamper-evident hash chain of all transactions.

Default cost model (configurable):
    cpu_ms:    0.001 IT per ms      (1 IT per second)
    io_bytes:  1e-9 IT per byte     (1 IT per GB)
    ram_bytes: 1e-10 IT per byte    (0.1 IT per GB-second)
"""

from __future__ import annotations

import hashlib
import json
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple


@dataclass
class LedgerEntry:
    """
    Single entry in the IT Token Ledger.

    Each entry records a debit or credit with full provenance.
    Entries form a hash chain for tamper-evidence.
    """
    # Identity
    entry_id: str = ""                 # Unique entry ID
    timestamp_utc: str = ""

    # Transaction
    syscall_id: str = ""               # Which syscall
    actor: str = ""                    # Who (actor/process)
    auth_id: str = ""                  # Pre-authorization ID (for matching pre_authorize → commit)

    # Cost model version
    cost_model_version: str = "v1"

    # Debits and credits
    debits: Dict[str, float] = field(default_factory=dict)   # {cpu_ms: X, io_bytes: Y, ram_bytes: Z}
    credits: Dict[str, float] = field(default_factory=dict)  # Rollback credits
    net_it_tokens: float = 0.0

    # Hash chain
    prev_hash: str = ""                # Hash of previous entry
    entry_hash: str = ""               # Hash of this entry (tamper-evident)

    def __post_init__(self):
        """Initialize timestamp if not provided."""
        if not self.timestamp_utc:
            self.timestamp_utc = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")

    def compute_hash(self) -> str:
        """Compute hash of this entry (excluding entry_hash)."""
        canonical = {
            "entry_id": self.entry_id,
            "timestamp_utc": self.timestamp_utc,
            "syscall_id": self.syscall_id,
            "actor": self.actor,
            "auth_id": self.auth_id,
            "cost_model_version": self.cost_model_version,
            "debits": {k: round(v, 9) for k, v in self.debits.items()},
            "credits": {k: round(v, 9) for k, v in self.credits.items()},
            "net_it_tokens": round(self.net_it_tokens, 9),
            "prev_hash": self.prev_hash,
        }
        blob = json.dumps(canonical, sort_keys=True, separators=(",", ":")).encode()
        return hashlib.sha256(blob).hexdigest()

    def finalize(self) -> "LedgerEntry":
        """Finalize entry by computing hash."""
        self.entry_hash = self.compute_hash()
        return self

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dict for JSON serialization."""
        return {
            "entry_id": self.entry_id,
            "timestamp_utc": self.timestamp_utc,
            "syscall_id": self.syscall_id,
            "actor": self.actor,
            "auth_id": self.auth_id,
            "cost_model_version": self.cost_model_version,
            "debits": self.debits,
            "credits": self.credits,
            "net_it_tokens": round(self.net_it_tokens, 9),
            "prev_hash": self.prev_hash,
            "entry_hash": self.entry_hash,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "LedgerEntry":
        """Create from dict."""
        return cls(
            entry_id=d.get("entry_id", ""),
            timestamp_utc=d.get("timestamp_utc", ""),
            syscall_id=d.get("syscall_id", ""),
            actor=d.get("actor", ""),
            auth_id=d.get("auth_id", ""),
            cost_model_version=d.get("cost_model_version", "v1"),
            debits=d.get("debits", {}),
            credits=d.get("credits", {}),
            net_it_tokens=d.get("net_it_tokens", 0.0),
            prev_hash=d.get("prev_hash", ""),
            entry_hash=d.get("entry_hash", ""),
        )


@dataclass
class PreAuthorization:
    """Pre-authorization hold for a syscall."""
    auth_id: str
    actor: str
    syscall_id: str
    estimated_cost: float
    timestamp_utc: str
    expires_at: float  # time.time() + timeout


class ITTokenLedger:
    """
    IT Token Ledger - Unified resource accounting.

    Features:
    - Pre-authorize → commit/rollback workflow
    - Configurable cost model
    - Tamper-evident hash chain
    - Per-actor balance tracking

    Usage:
        ledger = ITTokenLedger()
        ledger.set_balance("actor1", 1000.0)

        ok, err = ledger.pre_authorize("actor1", "cdna.decode.v1", 5.0)
        if ok:
            # ... do work ...
            entry = ledger.commit("actor1", "cdna.decode.v1", actual_cost, auth_id)
    """

    # Default cost model (IT tokens per unit)
    # All costs are computed deterministically with 9 decimal precision
    DEFAULT_COST_MODEL = {
        "syscall_base": 0.001,    # Base cost per syscall invocation
        "cpu_ms": 0.001,          # 1 IT token per second
        "io_bytes": 1e-9,         # 1 IT token per GB
        "ram_bytes_ms": 1e-13,    # 0.1 IT per GB-second RAM (ram_bytes * duration_ms)
    }

    COST_MODEL_VERSION = "v2"     # Bumped for syscall_base + ram_bytes_ms change
    PRE_AUTH_TIMEOUT_S = 300.0    # 5 minutes
    COST_PRECISION = 9            # Decimal places for deterministic costs

    def __init__(
        self,
        cost_model: Optional[Dict[str, float]] = None,
        default_balance: float = 1000.0,
    ):
        """
        Initialize ledger.

        Args:
            cost_model: Override default cost model
            default_balance: Default IT balance for new actors
        """
        self.cost_model = cost_model or self.DEFAULT_COST_MODEL.copy()
        self.default_balance = default_balance

        # State
        self._balances: Dict[str, float] = {}       # actor → balance
        self._entries: List[LedgerEntry] = []        # Hash chain
        self._pre_auths: Dict[str, PreAuthorization] = {}  # auth_id → PreAuthorization
        self._entry_counter: int = 0

    def _generate_auth_id(self) -> str:
        """Generate unique authorization ID."""
        import uuid
        return f"auth_{uuid.uuid4().hex[:12]}"

    def _generate_entry_id(self) -> str:
        """Generate unique entry ID."""
        self._entry_counter += 1
        return f"entry_{self._entry_counter:08d}"

    def _get_prev_hash(self) -> str:
        """Get hash of last entry (or genesis)."""
        if not self._entries:
            return hashlib.sha256(b"genesis").hexdigest()
        return self._entries[-1].entry_hash

    def compute_cost(self, breakdown: Dict[str, float]) -> float:
        """
        Compute total IT cost from breakdown using unified cost formula.

        Cost formula:
            total = syscall_base
                  + cpu_ms * cost_model["cpu_ms"]
                  + io_bytes * cost_model["io_bytes"]
                  + ram_bytes * duration_ms * cost_model["ram_bytes_ms"]

        All costs are rounded to 9 decimal places for determinism.

        Args:
            breakdown: {
                cpu_ms: CPU time in milliseconds,
                io_bytes: I/O bytes read/written,
                ram_bytes: Peak RAM usage in bytes,
                duration_ms: Duration in milliseconds (for RAM cost)
            }

        Returns:
            Total IT tokens (rounded to 9 decimal places)
        """
        # Start with base syscall cost
        total = self.cost_model.get("syscall_base", 0.001)

        # CPU time cost
        cpu_ms = breakdown.get("cpu_ms", 0.0)
        total += cpu_ms * self.cost_model.get("cpu_ms", 0.001)

        # I/O cost
        io_bytes = breakdown.get("io_bytes", 0.0)
        total += io_bytes * self.cost_model.get("io_bytes", 1e-9)

        # RAM cost (proportional to RAM * time)
        ram_bytes = breakdown.get("ram_bytes", 0.0)
        duration_ms = breakdown.get("duration_ms", breakdown.get("cpu_ms", 0.0))
        ram_cost = ram_bytes * duration_ms * self.cost_model.get("ram_bytes_ms", 1e-13)
        total += ram_cost

        # Round to COST_PRECISION decimal places for determinism
        return round(total, self.COST_PRECISION)

    def set_balance(self, actor: str, balance: float) -> None:
        """Set balance for an actor."""
        self._balances[actor] = balance

    def get_balance(self, actor: str) -> float:
        """Get balance for an actor (creates with default if missing)."""
        if actor not in self._balances:
            self._balances[actor] = self.default_balance
        return self._balances[actor]

    def pre_authorize(
        self,
        actor: str,
        syscall_id: str,
        estimated_cost: float,
    ) -> Tuple[bool, str, str]:
        """
        Pre-authorize IT spend for a syscall.

        Checks balance and holds estimated cost.

        Args:
            actor: Who is spending
            syscall_id: Which syscall
            estimated_cost: Estimated IT cost

        Returns:
            (success: bool, error: str, auth_id: str)
        """
        balance = self.get_balance(actor)

        # Check available (balance minus pending holds)
        pending = sum(
            pa.estimated_cost for pa in self._pre_auths.values()
            if pa.actor == actor
        )
        available = balance - pending

        if estimated_cost > available:
            return (
                False,
                f"ERESOURCE: insufficient IT tokens (need {estimated_cost:.3f}, have {available:.3f})",
                "",
            )

        # Create pre-authorization
        auth_id = self._generate_auth_id()
        self._pre_auths[auth_id] = PreAuthorization(
            auth_id=auth_id,
            actor=actor,
            syscall_id=syscall_id,
            estimated_cost=estimated_cost,
            timestamp_utc=datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
            expires_at=time.time() + self.PRE_AUTH_TIMEOUT_S,
        )

        return (True, "", auth_id)

    def commit(
        self,
        actor: str,
        syscall_id: str,
        actual_cost_breakdown: Dict[str, float],
        auth_id: str,
    ) -> LedgerEntry:
        """
        Commit a syscall cost to the ledger.

        Consumes the pre-authorization and records actual cost.

        Args:
            actor: Who is spending
            syscall_id: Which syscall
            actual_cost_breakdown: {cpu_ms: X, io_bytes: Y, ram_bytes: Z}
            auth_id: Pre-authorization ID

        Returns:
            LedgerEntry for the transaction
        """
        # Remove pre-authorization
        pre_auth = self._pre_auths.pop(auth_id, None)

        # Compute actual cost
        actual_cost = self.compute_cost(actual_cost_breakdown)

        # Debit balance
        self._balances[actor] = self.get_balance(actor) - actual_cost

        # Create ledger entry
        entry = LedgerEntry(
            entry_id=self._generate_entry_id(),
            syscall_id=syscall_id,
            actor=actor,
            auth_id=auth_id,
            cost_model_version=self.COST_MODEL_VERSION,
            debits=actual_cost_breakdown,
            net_it_tokens=-actual_cost,
            prev_hash=self._get_prev_hash(),
        )
        entry.finalize()

        self._entries.append(entry)
        return entry

    def rollback(self, auth_id: str) -> bool:
        """
        Rollback a pre-authorization (syscall failed or was cancelled).

        Args:
            auth_id: Pre-authorization ID

        Returns:
            True if rollback succeeded
        """
        if auth_id in self._pre_auths:
            del self._pre_auths[auth_id]
            return True
        return False

    def credit(
        self,
        actor: str,
        amount: float,
        reason: str = "manual_credit",
    ) -> LedgerEntry:
        """
        Credit IT tokens to an actor.

        Args:
            actor: Who to credit
            amount: Amount to credit
            reason: Reason for credit

        Returns:
            LedgerEntry for the transaction
        """
        self._balances[actor] = self.get_balance(actor) + amount

        entry = LedgerEntry(
            entry_id=self._generate_entry_id(),
            syscall_id=reason,
            actor=actor,
            auth_id="",
            cost_model_version=self.COST_MODEL_VERSION,
            credits={"credit": amount},
            net_it_tokens=amount,
            prev_hash=self._get_prev_hash(),
        )
        entry.finalize()

        self._entries.append(entry)
        return entry

    def verify_chain(self) -> Tuple[bool, int]:
        """
        Verify the ledger hash chain integrity.

        Returns:
            (valid: bool, bad_index: int or -1 if valid)
        """
        if not self._entries:
            return (True, -1)

        expected_prev = hashlib.sha256(b"genesis").hexdigest()

        for i, entry in enumerate(self._entries):
            # Check prev_hash links correctly
            if entry.prev_hash != expected_prev:
                return (False, i)

            # Check entry_hash is correct
            if entry.entry_hash != entry.compute_hash():
                return (False, i)

            expected_prev = entry.entry_hash

        return (True, -1)

    def get_entries(self) -> List[LedgerEntry]:
        """Get all ledger entries."""
        return self._entries.copy()

    def get_entries_for_actor(self, actor: str) -> List[LedgerEntry]:
        """Get ledger entries for a specific actor."""
        return [e for e in self._entries if e.actor == actor]

    def cleanup_expired_preauths(self) -> int:
        """Remove expired pre-authorizations. Returns count removed."""
        now = time.time()
        expired = [
            auth_id for auth_id, pa in self._pre_auths.items()
            if pa.expires_at < now
        ]
        for auth_id in expired:
            del self._pre_auths[auth_id]
        return len(expired)

    def stats(self) -> Dict[str, Any]:
        """Get ledger statistics."""
        return {
            "entry_count": len(self._entries),
            "actor_count": len(self._balances),
            "pending_preauths": len(self._pre_auths),
            "cost_model_version": self.COST_MODEL_VERSION,
            "cost_model": self.cost_model,
            "chain_valid": self.verify_chain()[0],
        }

    def to_dict(self) -> Dict[str, Any]:
        """Export ledger state for persistence."""
        return {
            "cost_model": self.cost_model,
            "default_balance": self.default_balance,
            "balances": self._balances,
            "entries": [e.to_dict() for e in self._entries],
            "entry_counter": self._entry_counter,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ITTokenLedger":
        """Restore ledger from dict."""
        ledger = cls(
            cost_model=d.get("cost_model"),
            default_balance=d.get("default_balance", 1000.0),
        )
        ledger._balances = d.get("balances", {})
        ledger._entries = [LedgerEntry.from_dict(e) for e in d.get("entries", [])]
        ledger._entry_counter = d.get("entry_counter", 0)
        return ledger
