# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Genesis Block & Provenance Chain

Implements the genesis block system for professional AI agent identity
verification (patent: George Scott Foley). Provides:

- Dataclasses for the five genesis block components (patent claim 3)
- RSA-2048 signing and verification (patent compliance)
- Ed25519 optional alternative for subsequent chain blocks
- SQLite-backed append-only chain storage

The genesis block anchors a Familiar instance's identity and provenance
to an immutable, cryptographically signed blockchain. SQLite serves as
the demo chain store (production target: Ethereum L2 / Arbitrum).

Dependencies: cryptography (already in requirements.txt)
"""

import hashlib
import json
import logging
import sqlite3
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa

logger = logging.getLogger(__name__)


# ── Genesis block component dataclasses (patent claim 3 a-e) ──


@dataclass
class ProfessionalCredentialAnchor:
    """(a) PCA — anchors professional credentials into the genesis block."""

    credential_hash: str
    authority_signatures: List[str] = field(default_factory=list)
    competency_metadata: Dict = field(default_factory=dict)
    regulatory_attestations: List[str] = field(default_factory=list)
    orcid: str = ""


@dataclass
class ProfessionalTemplateDerivation:
    """(b) PTD — records template origin and customization."""

    template_source_id: str = ""
    customization_params: Dict = field(default_factory=dict)
    individualization_sigs: List[str] = field(default_factory=list)
    evolution_provenance: List[str] = field(default_factory=list)


@dataclass
class VerifiedCapabilityAttestation:
    """(c) VCA — attests to verified professional capabilities."""

    expertise_domains: Dict[str, List[str]] = field(default_factory=dict)
    quality_thresholds: Dict[str, float] = field(default_factory=dict)
    standard_refs: List[str] = field(default_factory=list)
    validation_proofs: List[str] = field(default_factory=list)


@dataclass
class ProfessionalOwnershipFramework:
    """(d) POF — defines ownership, delegation, and succession."""

    multi_sig_control: Dict[str, str] = field(default_factory=dict)
    operator_authorization: Dict = field(default_factory=dict)
    succession_params: Dict = field(default_factory=dict)
    delegation_matrix: Dict[str, List[str]] = field(default_factory=dict)


@dataclass
class ProvenanceChainInit:
    """(e) PCI — initializes the provenance chain with creation metadata."""

    creation_timestamp: str = ""
    witness_proofs: List[str] = field(default_factory=list)
    authorization_sigs: List[str] = field(default_factory=list)
    audit_anchors: List[str] = field(default_factory=list)


# ── Block dataclasses ──


@dataclass
class GenesisBlock:
    """
    Genesis block (block 0) — container for the five patent components
    plus an RSA-2048 signature binding them together.
    """

    pca: ProfessionalCredentialAnchor
    ptd: ProfessionalTemplateDerivation
    vca: VerifiedCapabilityAttestation
    pof: ProfessionalOwnershipFramework
    pci: ProvenanceChainInit
    version: int = 1
    block_hash: str = ""
    signature: str = ""

    def compute_hash(self) -> str:
        """SHA-256 of canonical JSON (excludes block_hash and signature)."""
        d = self.to_dict()
        d.pop("block_hash", None)
        d.pop("signature", None)
        canonical = json.dumps(d, sort_keys=True, separators=(",", ":"))
        self.block_hash = hashlib.sha256(canonical.encode()).hexdigest()
        return self.block_hash

    def sign(self, private_key) -> str:
        """RSA-2048 sign the block hash. Computes hash first if empty."""
        if not self.block_hash:
            self.compute_hash()
        sig_bytes = private_key.sign(
            self.block_hash.encode(),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH,
            ),
            hashes.SHA256(),
        )
        self.signature = sig_bytes.hex()
        return self.signature

    def verify(self, public_key) -> bool:
        """Verify the RSA signature against the block hash."""
        if not self.block_hash or not self.signature:
            return False
        try:
            public_key.verify(
                bytes.fromhex(self.signature),
                self.block_hash.encode(),
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
            return True
        except Exception:
            return False

    def to_dict(self) -> dict:
        return {
            "version": self.version,
            "block_hash": self.block_hash,
            "signature": self.signature,
            "pca": asdict(self.pca),
            "ptd": asdict(self.ptd),
            "vca": asdict(self.vca),
            "pof": asdict(self.pof),
            "pci": asdict(self.pci),
        }

    @classmethod
    def from_dict(cls, d: dict) -> "GenesisBlock":
        return cls(
            version=d.get("version", 1),
            block_hash=d.get("block_hash", ""),
            signature=d.get("signature", ""),
            pca=ProfessionalCredentialAnchor(**d["pca"]),
            ptd=ProfessionalTemplateDerivation(**d["ptd"]),
            vca=VerifiedCapabilityAttestation(**d["vca"]),
            pof=ProfessionalOwnershipFramework(**d["pof"]),
            pci=ProvenanceChainInit(**d["pci"]),
        )


@dataclass
class ChainBlock:
    """
    A block in the provenance chain. Genesis is index 0;
    subsequent blocks carry provenance, export, or training payloads.
    """

    index: int
    timestamp: str
    prev_hash: str
    block_hash: str
    block_type: str  # "genesis" | "provenance" | "export" | "training"
    payload: Dict
    signature: str = ""

    def compute_hash(self) -> str:
        """SHA-256 of canonical JSON (excludes block_hash and signature)."""
        d = {
            "index": self.index,
            "timestamp": self.timestamp,
            "prev_hash": self.prev_hash,
            "block_type": self.block_type,
            "payload": self.payload,
        }
        canonical = json.dumps(d, sort_keys=True, separators=(",", ":"))
        self.block_hash = hashlib.sha256(canonical.encode()).hexdigest()
        return self.block_hash

    def sign(self, private_key, algorithm: str = "rsa") -> str:
        """Sign the block hash with RSA-2048 or Ed25519."""
        if not self.block_hash:
            self.compute_hash()
        if algorithm == "ed25519":
            sig_bytes = private_key.sign(self.block_hash.encode())
        else:
            sig_bytes = private_key.sign(
                self.block_hash.encode(),
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH,
                ),
                hashes.SHA256(),
            )
        self.signature = sig_bytes.hex()
        return self.signature

    def verify(self, public_key, algorithm: str = "rsa") -> bool:
        """Verify signature against the block hash."""
        if not self.block_hash or not self.signature:
            return False
        try:
            if algorithm == "ed25519":
                public_key.verify(
                    bytes.fromhex(self.signature),
                    self.block_hash.encode(),
                )
            else:
                public_key.verify(
                    bytes.fromhex(self.signature),
                    self.block_hash.encode(),
                    padding.PSS(
                        mgf=padding.MGF1(hashes.SHA256()),
                        salt_length=padding.PSS.MAX_LENGTH,
                    ),
                    hashes.SHA256(),
                )
            return True
        except Exception:
            return False


# ── RSA key management ──


def generate_keypair() -> Tuple[bytes, bytes]:
    """
    Generate an RSA-2048 keypair.

    Returns (private_pem, public_pem) as PEM-encoded bytes.
    """
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
    )
    private_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    public_pem = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    return private_pem, public_pem


def load_or_create_keypair(key_dir: Path):
    """
    Load RSA keypair from key_dir, or create one if it doesn't exist.

    Private key is encrypted at rest if FAMILIAR_ENCRYPTION_KEY is set.

    Returns (private_key, public_key) as cryptography key objects.
    """
    key_dir.mkdir(parents=True, exist_ok=True)
    priv_path = key_dir / "chain_private.pem"
    pub_path = key_dir / "chain_public.pem"

    if priv_path.exists() and pub_path.exists():
        priv_pem = priv_path.read_bytes()
        # Decrypt private key if it was encrypted with Fernet
        if priv_pem.startswith(b"gAAAAA"):
            try:
                from familiar.core.encryption import EncryptedStorage

                storage = EncryptedStorage()
                if storage.is_available:
                    priv_pem = storage.decrypt(priv_pem.decode()).encode()
            except Exception:
                pass  # Fall through — key may be plaintext despite prefix
        private_key = serialization.load_pem_private_key(priv_pem, password=None)
        public_key = serialization.load_pem_public_key(pub_path.read_bytes())
        return private_key, public_key

    # Generate new keypair
    private_pem, public_pem = generate_keypair()

    # Encrypt private key at rest if encryption is available
    stored_priv = private_pem
    try:
        from familiar.core.encryption import EncryptedStorage

        storage = EncryptedStorage()
        if storage.is_available:
            stored_priv = storage.encrypt(private_pem.decode()).encode()
    except Exception:
        pass

    priv_path.write_bytes(stored_priv)
    priv_path.chmod(0o600)
    pub_path.write_bytes(public_pem)

    private_key = serialization.load_pem_private_key(private_pem, password=None)
    public_key = serialization.load_pem_public_key(public_pem)
    logger.info("Generated new RSA-2048 chain keypair at %s", key_dir)
    return private_key, public_key


# ── SQLite-backed append-only chain store ──

_CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS chain (
    idx INTEGER PRIMARY KEY,
    timestamp TEXT NOT NULL,
    prev_hash TEXT NOT NULL,
    block_hash TEXT NOT NULL UNIQUE,
    block_type TEXT NOT NULL,
    payload TEXT NOT NULL,
    signature TEXT NOT NULL
);
"""


class ChainStore:
    """SQLite-backed append-only provenance chain."""

    def __init__(self, db_path: Path):
        self.db_path = db_path
        db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()

    def _init_db(self):
        with sqlite3.connect(str(self.db_path)) as conn:
            conn.execute(_CREATE_TABLE_SQL)

    def append(self, block: ChainBlock) -> None:
        """Append a block to the chain. Validates index continuity."""
        with sqlite3.connect(str(self.db_path)) as conn:
            current_len = conn.execute("SELECT COUNT(*) FROM chain").fetchone()[0]
            if block.index != current_len:
                raise ValueError(f"Block index {block.index} != expected {current_len}")
            conn.execute(
                "INSERT INTO chain (idx, timestamp, prev_hash, block_hash, "
                "block_type, payload, signature) VALUES (?, ?, ?, ?, ?, ?, ?)",
                (
                    block.index,
                    block.timestamp,
                    block.prev_hash,
                    block.block_hash,
                    block.block_type,
                    json.dumps(block.payload, default=str),
                    block.signature,
                ),
            )

    def _row_to_block(self, row) -> ChainBlock:
        return ChainBlock(
            index=row[0],
            timestamp=row[1],
            prev_hash=row[2],
            block_hash=row[3],
            block_type=row[4],
            payload=json.loads(row[5]),
            signature=row[6],
        )

    def get_genesis(self) -> Optional[ChainBlock]:
        """Return the genesis block (index 0), or None."""
        return self.get_block(0)

    def get_latest(self) -> Optional[ChainBlock]:
        """Return the latest block in the chain."""
        with sqlite3.connect(str(self.db_path)) as conn:
            row = conn.execute(
                "SELECT idx, timestamp, prev_hash, block_hash, block_type, "
                "payload, signature FROM chain ORDER BY idx DESC LIMIT 1"
            ).fetchone()
        if row is None:
            return None
        return self._row_to_block(row)

    def get_block(self, index: int) -> Optional[ChainBlock]:
        """Return a specific block by index."""
        with sqlite3.connect(str(self.db_path)) as conn:
            row = conn.execute(
                "SELECT idx, timestamp, prev_hash, block_hash, block_type, "
                "payload, signature FROM chain WHERE idx = ?",
                (index,),
            ).fetchone()
        if row is None:
            return None
        return self._row_to_block(row)

    def get_chain(self, start: int = 0, limit: int = 100) -> List[ChainBlock]:
        """Return a slice of the chain."""
        with sqlite3.connect(str(self.db_path)) as conn:
            rows = conn.execute(
                "SELECT idx, timestamp, prev_hash, block_hash, block_type, "
                "payload, signature FROM chain WHERE idx >= ? "
                "ORDER BY idx ASC LIMIT ?",
                (start, limit),
            ).fetchall()
        return [self._row_to_block(r) for r in rows]

    def verify_chain(self) -> Tuple[bool, Optional[str]]:
        """
        Verify chain integrity: hash linking and hash recomputation.

        Returns (valid, error_message). error_message is None if valid.
        """
        blocks = self.get_chain(start=0, limit=999999)
        if not blocks:
            return True, None

        for i, block in enumerate(blocks):
            # Verify index continuity
            if block.index != i:
                return False, f"Block {i}: expected index {i}, got {block.index}"

            # Verify hash links
            if i == 0:
                if block.prev_hash != "0" * 64:
                    return False, "Genesis block prev_hash is not the null hash"
            else:
                if block.prev_hash != blocks[i - 1].block_hash:
                    return False, (
                        f"Block {i}: prev_hash mismatch "
                        f"(expected {blocks[i - 1].block_hash[:16]}..., "
                        f"got {block.prev_hash[:16]}...)"
                    )

            # Recompute and verify block hash
            saved_hash = block.block_hash
            block.compute_hash()
            if block.block_hash != saved_hash:
                return False, (
                    f"Block {i}: hash mismatch "
                    f"(stored {saved_hash[:16]}..., "
                    f"computed {block.block_hash[:16]}...)"
                )

        return True, None

    def chain_length(self) -> int:
        """Return number of blocks in the chain."""
        with sqlite3.connect(str(self.db_path)) as conn:
            return conn.execute("SELECT COUNT(*) FROM chain").fetchone()[0]
