#!/usr/bin/env python3
"""
transaction.py — UTXO Model & Transaction Format
Montana Protocol v3.2

UTXO (Unspent Transaction Output) — модель данных Таймчейна.
Балансов нет. Есть непотраченные выходы.

Каждый UTXO — запись: «на адрес mt… можно потратить N Ɉ,
если предъявишь подпись от ключа этого адреса».

Post-quantum security: ML-DSA-65 (FIPS 204) signatures.
"""

import hashlib
import json
import sqlite3
import time
import threading
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Dict

from node_crypto import (
    sign_message,
    verify_signature,
    public_key_to_address,
    validate_address,
)


# ═══════════════════════════════════════════════════════════════════════════════
#                         TRANSACTION FORMAT
# ═══════════════════════════════════════════════════════════════════════════════


@dataclass
class TxInput:
    """
    Вход транзакции — ссылка на непотраченный выход.

    При трате UTXO раскрывается публичный ключ владельца
    и предоставляется ML-DSA-65 подпись tx_hash.
    """
    tx_hash: str          # Хеш транзакции, создавшей выход
    output_idx: int       # Индекс выхода в той транзакции (0, 1, ...)
    pubkey: str           # Публичный ключ отправителя (1952 bytes hex)
    signature: str        # ML-DSA-65 подпись tx_hash (3309 bytes hex)

    def to_dict(self) -> Dict:
        return {
            "tx_hash": self.tx_hash,
            "output_idx": self.output_idx,
            "pubkey": self.pubkey,
            "signature": self.signature,
        }

    def to_dict_unsigned(self) -> Dict:
        """Для вычисления tx_hash (без подписей)"""
        return {
            "tx_hash": self.tx_hash,
            "output_idx": self.output_idx,
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "TxInput":
        return cls(
            tx_hash=d["tx_hash"],
            output_idx=d["output_idx"],
            pubkey=d.get("pubkey", ""),
            signature=d.get("signature", ""),
        )


@dataclass
class TxOutput:
    """
    Выход транзакции — новый UTXO.

    Адрес получателя и сумма в Ɉ (целое число).
    """
    address: str          # mt... адрес получателя (42 символа)
    amount: int           # Сумма в Ɉ (целое, без дробей)

    def to_dict(self) -> Dict:
        return {
            "address": self.address,
            "amount": self.amount,
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "TxOutput":
        return cls(address=d["address"], amount=d["amount"])


@dataclass
class Transaction:
    """
    Транзакция Montana Protocol.

    Два типа:
    - coinbase: нет входов, создаёт новые Ɉ (эмиссия в τ₂)
    - transfer: тратит UTXO, создаёт новые выходы

    tx_hash вычисляется из canonical serialization
    (inputs без подписей + outputs + timestamp + tx_type).
    """
    inputs: List[TxInput]
    outputs: List[TxOutput]
    timestamp: int            # Наносекунды UTC
    tx_type: str              # "coinbase" | "transfer"
    tx_hash: str = ""         # Вычисляется автоматически
    nonce: str = ""           # Gemini fix #6: uniqueness for coinbase (t2_index)

    def compute_hash(self) -> str:
        """
        Каноническая сериализация → SHA-256.
        Inputs без подписей (чтобы хеш не зависел от подписей).
        Gemini fix #6: nonce included for coinbase uniqueness.
        """
        data = {
            "inputs": [inp.to_dict_unsigned() for inp in self.inputs],
            "outputs": [out.to_dict() for out in self.outputs],
            "timestamp": self.timestamp,
            "tx_type": self.tx_type,
        }
        if hasattr(self, 'nonce') and self.nonce:
            data["nonce"] = self.nonce
        canonical = json.dumps(data, sort_keys=True, ensure_ascii=True)
        return hashlib.sha256(canonical.encode("utf-8")).hexdigest()

    def to_dict(self) -> Dict:
        d = {
            "tx_hash": self.tx_hash,
            "inputs": [inp.to_dict() for inp in self.inputs],
            "outputs": [out.to_dict() for out in self.outputs],
            "timestamp": self.timestamp,
            "tx_type": self.tx_type,
        }
        if self.nonce:
            d["nonce"] = self.nonce
        return d

    @classmethod
    def from_dict(cls, d: Dict) -> "Transaction":
        tx = cls(
            inputs=[TxInput.from_dict(i) for i in d.get("inputs", [])],
            outputs=[TxOutput.from_dict(o) for o in d.get("outputs", [])],
            timestamp=d["timestamp"],
            tx_type=d["tx_type"],
            tx_hash=d.get("tx_hash", ""),
            nonce=d.get("nonce", ""),
        )
        if not tx.tx_hash:
            tx.tx_hash = tx.compute_hash()
        return tx


# ═══════════════════════════════════════════════════════════════════════════════
#                         TRANSACTION CREATION
# ═══════════════════════════════════════════════════════════════════════════════


# Gemini R3 fix #2: SQLite INTEGER safety (max signed 64-bit)
MAX_SAFE_AMOUNT = 4_611_686_018_427_387_903  # 2^62 - 1


def create_coinbase_tx(address: str, amount: int, t2_index: int) -> Transaction:
    """
    Создаёт coinbase-транзакцию (эмиссия).

    Нет входов. Один выход на адрес участника.
    Количество Ɉ = секунды присутствия × halving_coefficient.

    Args:
        address: mt... адрес получателя
        amount: количество Ɉ
        t2_index: номер τ₂ блока (для уникальности хеша)

    Returns:
        Transaction с tx_type="coinbase"
    """
    if amount <= 0:
        raise ValueError(f"Coinbase amount must be positive: {amount}")
    if amount > MAX_SAFE_AMOUNT:
        raise ValueError(f"Coinbase amount exceeds safe limit: {amount} > {MAX_SAFE_AMOUNT}")
    if not validate_address(address):
        raise ValueError(f"Invalid address: {address}")

    tx = Transaction(
        inputs=[],
        outputs=[TxOutput(address=address, amount=amount)],
        timestamp=time.time_ns(),
        tx_type="coinbase",
        nonce=f"coinbase:{t2_index}:{address}",  # Gemini fix #6: unique per t2_index
    )
    tx.tx_hash = tx.compute_hash()
    return tx


def create_transfer_tx(
    utxos_to_spend: List[Tuple[str, int, int]],  # [(tx_hash, output_idx, amount), ...]
    recipient: str,
    amount: int,
    change_address: str,
    private_key: str,
    public_key: str,
) -> Transaction:
    """
    Создаёт transfer-транзакцию.

    1. Собирает inputs из UTXO
    2. Создаёт outputs: получатель + сдача
    3. Вычисляет tx_hash
    4. Подписывает каждый input

    Args:
        utxos_to_spend: список UTXO для траты [(tx_hash, output_idx, amount)]
        recipient: mt... адрес получателя
        amount: сколько отправить
        change_address: mt... адрес для сдачи (обычно = отправитель)
        private_key: hex приватный ключ ML-DSA-65
        public_key: hex публичный ключ ML-DSA-65

    Returns:
        Подписанная Transaction

    Raises:
        ValueError: если сумма inputs < amount или адреса невалидны
    """
    if not validate_address(recipient):
        raise ValueError(f"Invalid recipient address: {recipient}")
    if not validate_address(change_address):
        raise ValueError(f"Invalid change address: {change_address}")
    if amount <= 0:
        raise ValueError(f"Amount must be positive: {amount}")
    if amount > MAX_SAFE_AMOUNT:
        raise ValueError(f"Transfer amount exceeds safe limit: {amount} > {MAX_SAFE_AMOUNT}")

    total_input = sum(u[2] for u in utxos_to_spend)
    if total_input < amount:
        raise ValueError(
            f"Insufficient funds: have {total_input}, need {amount}"
        )

    # Создаём inputs (пока без подписей)
    inputs = [
        TxInput(
            tx_hash=u[0],
            output_idx=u[1],
            pubkey=public_key,
            signature="",  # Заполним после вычисления tx_hash
        )
        for u in utxos_to_spend
    ]

    # Создаём outputs
    outputs = [TxOutput(address=recipient, amount=amount)]
    change = total_input - amount
    if change > 0:
        outputs.append(TxOutput(address=change_address, amount=change))

    # Создаём транзакцию и вычисляем хеш
    tx = Transaction(
        inputs=inputs,
        outputs=outputs,
        timestamp=time.time_ns(),
        tx_type="transfer",
    )
    tx.tx_hash = tx.compute_hash()

    # Подписываем каждый input
    for inp in tx.inputs:
        inp.signature = sign_message(private_key, tx.tx_hash)

    return tx


# ═══════════════════════════════════════════════════════════════════════════════
#                         TRANSACTION VALIDATION
# ═══════════════════════════════════════════════════════════════════════════════


def validate_transaction(tx: Transaction, utxo_set: "UTXOSet") -> Tuple[bool, str]:
    """
    Полная валидация транзакции.

    Для coinbase:
        1. inputs пуст
        2. tx_type == "coinbase"
        3. Все outputs: amount > 0, address валиден

    Для transfer:
        1. Каждый input ссылается на существующий непотраченный UTXO
        2. address(input.pubkey) == UTXO.address
        3. ML_DSA_65.verify(pubkey, tx_hash, signature) == True
        4. sum(inputs) == sum(outputs) (Montana: без комиссий)
        5. Все outputs: amount > 0, address валиден
        6. Нет дубликатов inputs (double spend внутри одной tx)
        7. tx_hash корректен (пересчёт совпадает)

    Returns:
        (is_valid, error_message)
    """
    # Проверка tx_hash
    expected_hash = tx.compute_hash()
    if tx.tx_hash != expected_hash:
        return False, f"Invalid tx_hash: expected {expected_hash[:16]}..., got {tx.tx_hash[:16]}..."

    # Проверка outputs
    if not tx.outputs:
        return False, "Transaction has no outputs"
    for i, out in enumerate(tx.outputs):
        if out.amount <= 0:
            return False, f"Output #{i}: amount must be positive ({out.amount})"
        if out.amount > MAX_SAFE_AMOUNT:
            return False, f"Output #{i}: amount exceeds safe limit ({out.amount} > {MAX_SAFE_AMOUNT})"
        if not validate_address(out.address):
            return False, f"Output #{i}: invalid address {out.address}"

    if tx.tx_type == "coinbase":
        # Coinbase: нет входов
        if tx.inputs:
            return False, "Coinbase transaction must have no inputs"
        return True, "OK"

    elif tx.tx_type == "transfer":
        if not tx.inputs:
            return False, "Transfer transaction must have inputs"

        # Проверка дубликатов inputs
        seen_inputs = set()
        for inp in tx.inputs:
            key = (inp.tx_hash, inp.output_idx)
            if key in seen_inputs:
                return False, f"Duplicate input: {inp.tx_hash[:16]}...:{inp.output_idx}"
            seen_inputs.add(key)

        # ML-DSA-65 key sizes (hex)
        EXPECTED_PUBKEY_LEN = 3904     # 1952 bytes = 3904 hex chars
        EXPECTED_SIG_LEN = 6618        # 3309 bytes = 6618 hex chars

        # Проверка каждого input
        total_input = 0
        for i, inp in enumerate(tx.inputs):
            # 0. Pubkey format validation (ML-DSA-65: 1952 bytes = 3904 hex)
            if not inp.pubkey or len(inp.pubkey) != EXPECTED_PUBKEY_LEN:
                return False, (
                    f"Input #{i}: invalid pubkey length "
                    f"(expected {EXPECTED_PUBKEY_LEN}, got {len(inp.pubkey) if inp.pubkey else 0})"
                )
            if not all(c in '0123456789abcdef' for c in inp.pubkey):
                return False, f"Input #{i}: pubkey contains non-hex characters"

            # 0b. Signature format validation
            if not inp.signature or len(inp.signature) != EXPECTED_SIG_LEN:
                return False, (
                    f"Input #{i}: invalid signature length "
                    f"(expected {EXPECTED_SIG_LEN}, got {len(inp.signature) if inp.signature else 0})"
                )

            # 1. UTXO существует и не потрачен
            utxo = utxo_set.get_utxo(inp.tx_hash, inp.output_idx)
            if utxo is None:
                return False, f"Input #{i}: UTXO not found ({inp.tx_hash[:16]}...:{inp.output_idx})"
            if utxo["spent_by_tx"] is not None:
                return False, f"Input #{i}: UTXO already spent by {utxo['spent_by_tx'][:16]}..."

            # 2. address(pubkey) == UTXO.address
            derived_address = public_key_to_address(inp.pubkey)
            if derived_address != utxo["address"]:
                return False, f"Input #{i}: pubkey address mismatch (derived={derived_address}, utxo={utxo['address']})"

            # 3. ML-DSA-65 подпись валидна
            if not verify_signature(inp.pubkey, tx.tx_hash, inp.signature):
                return False, f"Input #{i}: invalid ML-DSA-65 signature"

            total_input += utxo["amount"]

        # 4. sum(inputs) == sum(outputs)
        total_output = sum(out.amount for out in tx.outputs)
        if total_input != total_output:
            return False, f"Input/output mismatch: inputs={total_input}, outputs={total_output}"

        return True, "OK"

    else:
        return False, f"Unknown tx_type: {tx.tx_type}"


# ═══════════════════════════════════════════════════════════════════════════════
#                         UTXO SET (SQLite)
# ═══════════════════════════════════════════════════════════════════════════════


class UTXOSet:
    """
    Управление непотраченными выходами транзакций.

    Каждый UTXO — запись в SQLite:
    - (tx_hash, output_idx) — уникальный идентификатор
    - address — владелец
    - amount — сумма в Ɉ
    - created_in_tau1 — в каком τ₁ блоке создан
    - spent_by_tx — NULL если не потрачен, иначе tx_hash потратившей TX

    Баланс адреса = SUM(amount) WHERE address=? AND spent_by_tx IS NULL
    """

    def __init__(self, db_path: str):
        self.db_path = db_path
        self._lock = threading.Lock()
        self._init_db()

    def _init_db(self):
        """Создаём таблицу UTXO если не существует"""
        with sqlite3.connect(self.db_path) as conn:
            # IMMUTABILITY: WAL mode + crash protection
            conn.execute("PRAGMA journal_mode=WAL")
            conn.execute("PRAGMA synchronous=FULL")

            # IMMUTABILITY: integrity check on startup
            result = conn.execute("PRAGMA integrity_check").fetchone()
            if result[0] != "ok":
                raise RuntimeError(f"UTXO DATABASE CORRUPTION DETECTED: {result[0]}")

            conn.execute("""
                CREATE TABLE IF NOT EXISTS utxos (
                    tx_hash TEXT NOT NULL,
                    output_idx INTEGER NOT NULL,
                    address TEXT NOT NULL,
                    amount INTEGER NOT NULL,
                    created_in_tau1 INTEGER NOT NULL,
                    spent_by_tx TEXT,
                    PRIMARY KEY (tx_hash, output_idx)
                )
            """)
            # Индекс для быстрого запроса баланса (непотраченные UTXO адреса)
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_utxo_address_unspent
                ON utxos(address) WHERE spent_by_tx IS NULL
            """)
            # Индекс для проверки total supply
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_utxo_unspent
                ON utxos(spent_by_tx) WHERE spent_by_tx IS NULL
            """)
            conn.commit()

    def _conn(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def add_outputs(self, tx: Transaction, tau1_block_num: int):
        """
        Добавить выходы транзакции в UTXO set.
        Вызывается при включении TX в τ₁ блок.
        """
        with self._lock:
            with self._conn() as conn:
                for idx, out in enumerate(tx.outputs):
                    conn.execute(
                        """INSERT OR IGNORE INTO utxos
                           (tx_hash, output_idx, address, amount, created_in_tau1, spent_by_tx)
                           VALUES (?, ?, ?, ?, ?, NULL)""",
                        (tx.tx_hash, idx, out.address, out.amount, tau1_block_num),
                    )
                conn.commit()

    def spend_output(self, tx_hash: str, output_idx: int, spending_tx_hash: str) -> bool:
        """
        Пометить UTXO как потраченный.

        Returns:
            True если успешно, False если UTXO не найден или уже потрачен.
        """
        with self._lock:
            with self._conn() as conn:
                cursor = conn.execute(
                    """UPDATE utxos SET spent_by_tx = ?
                       WHERE tx_hash = ? AND output_idx = ? AND spent_by_tx IS NULL""",
                    (spending_tx_hash, tx_hash, output_idx),
                )
                conn.commit()
                return cursor.rowcount == 1

    def get_utxo(self, tx_hash: str, output_idx: int) -> Optional[Dict]:
        """Получить UTXO по идентификатору"""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT * FROM utxos WHERE tx_hash = ? AND output_idx = ?",
                (tx_hash, output_idx),
            ).fetchone()
            if row:
                return dict(row)
            return None

    def get_balance(self, address: str) -> int:
        """
        Баланс адреса = сумма непотраченных UTXO.
        """
        with self._conn() as conn:
            row = conn.execute(
                "SELECT COALESCE(SUM(amount), 0) as balance FROM utxos WHERE address = ? AND spent_by_tx IS NULL",
                (address,),
            ).fetchone()
            return row["balance"]

    def get_utxos_for_address(self, address: str) -> List[Dict]:
        """Все непотраченные UTXO адреса"""
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT * FROM utxos WHERE address = ? AND spent_by_tx IS NULL ORDER BY amount DESC",
                (address,),
            ).fetchall()
            return [dict(r) for r in rows]

    def is_unspent(self, tx_hash: str, output_idx: int) -> bool:
        """Проверить что UTXO не потрачен"""
        utxo = self.get_utxo(tx_hash, output_idx)
        return utxo is not None and utxo["spent_by_tx"] is None

    def total_supply(self) -> int:
        """
        Сумма непотраченных UTXO = текущее денежное предложение.
        В UTXO модели потраченные выходы уничтожены — считаются только живые.
        """
        with self._conn() as conn:
            row = conn.execute(
                "SELECT COALESCE(SUM(amount), 0) as total FROM utxos WHERE spent_by_tx IS NULL"
            ).fetchone()
            return row["total"]

    def total_unspent(self) -> int:
        """Alias для total_supply()"""
        return self.total_supply()

    def total_coinbase_emissions(self) -> int:
        """
        Сумма всех coinbase выходов (эмиссия с генезиса).
        Coinbase UTXO не имеют входов, поэтому их tx_hash уникальны
        и created_in_tau1 >= 0.
        Это НЕ зависит от того, потрачены они или нет.
        """
        with self._conn() as conn:
            row = conn.execute(
                "SELECT COALESCE(SUM(amount), 0) as total FROM utxos"
            ).fetchone()
            return row["total"]

    def utxo_count(self) -> int:
        """Количество непотраченных UTXO"""
        with self._conn() as conn:
            row = conn.execute(
                "SELECT COUNT(*) as cnt FROM utxos WHERE spent_by_tx IS NULL"
            ).fetchone()
            return row["cnt"]

    def all_balances(self) -> Dict[str, int]:
        """Все адреса с ненулевым UTXO-балансом"""
        with self._conn() as conn:
            rows = conn.execute(
                "SELECT address, SUM(amount) as balance FROM utxos "
                "WHERE spent_by_tx IS NULL GROUP BY address HAVING balance > 0"
            ).fetchall()
            return {r["address"]: r["balance"] for r in rows}

    def select_utxos_for_amount(self, address: str, amount: int) -> List[Dict]:
        """
        Выбрать UTXO для покрытия суммы (жадный алгоритм — крупные первыми).

        Returns:
            Список UTXO, сумма которых >= amount.
            Пустой список если недостаточно средств.
        """
        utxos = self.get_utxos_for_address(address)
        selected = []
        total = 0
        for utxo in utxos:
            selected.append(utxo)
            total += utxo["amount"]
            if total >= amount:
                return selected
        return []  # Недостаточно средств

    def apply_transaction(self, tx: Transaction, tau1_block_num: int) -> Tuple[bool, str]:
        """
        Атомарно применить транзакцию к UTXO set.

        Для coinbase: добавить outputs.
        Для transfer: пометить inputs как потраченные + добавить outputs.

        Returns:
            (success, error_message)
        """
        if tx.tx_type == "coinbase":
            self.add_outputs(tx, tau1_block_num)
            return True, "OK"

        elif tx.tx_type == "transfer":
            with self._lock:
                with self._conn() as conn:
                    # Проверяем и тратим inputs атомарно
                    for inp in tx.inputs:
                        cursor = conn.execute(
                            """UPDATE utxos SET spent_by_tx = ?
                               WHERE tx_hash = ? AND output_idx = ? AND spent_by_tx IS NULL""",
                            (tx.tx_hash, inp.tx_hash, inp.output_idx),
                        )
                        if cursor.rowcount != 1:
                            conn.rollback()
                            return False, f"Failed to spend UTXO {inp.tx_hash[:16]}...:{inp.output_idx}"

                    # Добавляем outputs
                    for idx, out in enumerate(tx.outputs):
                        conn.execute(
                            """INSERT INTO utxos
                               (tx_hash, output_idx, address, amount, created_in_tau1, spent_by_tx)
                               VALUES (?, ?, ?, ?, ?, NULL)""",
                            (tx.tx_hash, idx, out.address, out.amount, tau1_block_num),
                        )

                    conn.commit()
                    return True, "OK"

        return False, f"Unknown tx_type: {tx.tx_type}"

    def verify_supply_invariant(self) -> Tuple[bool, str]:
        """
        Критический инвариант: transfer не создаёт и не уничтожает Ɉ.

        Проверка: сумма непотраченных UTXO (денежное предложение) > 0
        и совпадает с разницей: (все coinbase выходы) - (потраченные coinbase).

        Упрощённо: total_supply >= 0 (не отрицательный баланс).
        Полная проверка: для каждой transfer TX, sum(inputs) == sum(outputs).
        """
        supply = self.total_supply()
        if supply < 0:
            return False, f"Negative supply: {supply}"

        # Проверяем что каждая transfer TX сбалансирована:
        # Все потраченные UTXO (spent_by_tx != NULL) группируем по spending TX
        # и проверяем что для каждой TX сумма потраченных == сумма созданных
        with self._conn() as conn:
            # Суммы потраченных по каждой spending TX
            spent_rows = conn.execute("""
                SELECT spent_by_tx, SUM(amount) as spent_amount
                FROM utxos
                WHERE spent_by_tx IS NOT NULL
                GROUP BY spent_by_tx
            """).fetchall()

            for row in spent_rows:
                spending_tx = row["spent_by_tx"]
                spent_amount = row["spent_amount"]

                # Сумма выходов этой TX
                created = conn.execute("""
                    SELECT COALESCE(SUM(amount), 0) as created_amount
                    FROM utxos
                    WHERE tx_hash = ?
                """, (spending_tx,)).fetchone()

                created_amount = created["created_amount"]
                if spent_amount != created_amount:
                    return False, (
                        f"Transfer TX {spending_tx[:16]}... imbalanced: "
                        f"spent={spent_amount}, created={created_amount}"
                    )

        return True, f"OK (supply={supply}Ɉ)"