#!/usr/bin/env python3
"""
LuckySt Syndicate - Report Encryption & Key Management

AES-256-GCM encryption for trading session reports.
Redis-based key storage with 24-hour access grants for pledged syndicate members.
"""

import os
import json
import secrets
from dataclasses import dataclass
from typing import Optional
from pathlib import Path

import yaml
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

# Wire format version
REPORT_VERSION = 1

REGISTRY_PATH = Path(__file__).parent / "registry.yaml"


# ---------------------------------------------------------------------------
# Encryption / Decryption
# ---------------------------------------------------------------------------

@dataclass
class EncryptedReport:
    """An AES-256-GCM encrypted report with its key material."""
    ciphertext: bytes
    nonce: bytes  # 12 bytes
    key: bytes    # 32 bytes

    def to_chain_bytes(self) -> bytes:
        """
        Serialize for on-chain storage.

        Format: [version:1][nonce:12][ciphertext:N]
        The AES key is NOT included — it's stored in Redis.
        """
        version_byte = REPORT_VERSION.to_bytes(1, "big")
        return version_byte + self.nonce + self.ciphertext

    @staticmethod
    def parse_chain_bytes(data: bytes) -> tuple[int, bytes, bytes]:
        """
        Parse on-chain bytes back into components.

        Returns: (version, nonce, ciphertext)
        """
        version = data[0]
        nonce = data[1:13]
        ciphertext = data[13:]
        return version, nonce, ciphertext


def encrypt_report(plaintext: str) -> EncryptedReport:
    """
    Encrypt a report string with AES-256-GCM.

    Generates a random 256-bit key and 96-bit nonce.
    Returns EncryptedReport with ciphertext, nonce, and key.
    """
    key = secrets.token_bytes(32)
    nonce = secrets.token_bytes(12)
    aesgcm = AESGCM(key)
    ciphertext = aesgcm.encrypt(nonce, plaintext.encode("utf-8"), None)
    return EncryptedReport(ciphertext=ciphertext, nonce=nonce, key=key)


def decrypt_report(ciphertext: bytes, nonce: bytes, key: bytes) -> str:
    """
    Decrypt an AES-256-GCM encrypted report.

    Args:
        ciphertext: The encrypted data
        nonce: 12-byte nonce used during encryption
        key: 32-byte AES key

    Returns:
        Decrypted plaintext string
    """
    aesgcm = AESGCM(key)
    plaintext = aesgcm.decrypt(nonce, ciphertext, None)
    return plaintext.decode("utf-8")


# ---------------------------------------------------------------------------
# Redis Key Management
# ---------------------------------------------------------------------------

class ReportKeyManager:
    """
    Manages AES keys for encrypted reports via Redis.

    - Authors keep keys indefinitely (no expiry)
    - Other pledged syndicate members can request 24h access
    - Pledge verification checks registry.yaml
    """

    # Redis key prefixes
    KEY_PREFIX = "report:key:"        # report:key:{tx_hash} -> JSON metadata
    GRANT_PREFIX = "report:grant:"    # report:grant:{tx_hash}:{address} -> AES key hex
    INDEX_PREFIX = "report:index:"    # report:index:{author} -> sorted set of tx hashes

    GRANT_TTL = 86400  # 24 hours

    def __init__(self, redis_client):
        self.redis = redis_client

    def store_report_key(
        self,
        tx_hash: str,
        aes_key: bytes,
        author_address: str,
        instance_id: int,
        markets: list[str],
        timestamp: str,
    ):
        """
        Store the AES key for a report. Called by the report author.

        The key is stored with no expiry — the author always has access.
        Also indexes the report under the author's address for listing.
        """
        metadata = {
            "aes_key_hex": aes_key.hex(),
            "author": author_address.lower(),
            "instance_id": instance_id,
            "markets": markets,
            "timestamp": timestamp,
            "tx_hash": tx_hash,
        }

        # Store key metadata (no expiry)
        self.redis.set(
            f"{self.KEY_PREFIX}{tx_hash}",
            json.dumps(metadata),
        )

        # Index under author for listing
        self.redis.zadd(
            f"{self.INDEX_PREFIX}{author_address.lower()}",
            {tx_hash: float(int(timestamp) if timestamp.isdigit() else 0)},
        )

    def grant_access(
        self,
        tx_hash: str,
        requester_address: str,
        author_address: str,
    ) -> Optional[str]:
        """
        Grant 24-hour access to a report's AES key.

        Verifies the requester is a pledged syndicate member.
        Returns the AES key hex string, or None if denied.
        """
        # Check if requester is the author (always allowed)
        if requester_address.lower() == author_address.lower():
            meta = self._get_metadata(tx_hash)
            return meta["aes_key_hex"] if meta else None

        # Verify requester is pledged
        if not self._is_pledged(requester_address):
            return None

        # Fetch the key
        meta = self._get_metadata(tx_hash)
        if not meta:
            return None

        # Verify author matches
        if meta["author"] != author_address.lower():
            return None

        aes_key_hex = meta["aes_key_hex"]

        # Store a time-limited grant
        grant_key = f"{self.GRANT_PREFIX}{tx_hash}:{requester_address.lower()}"
        self.redis.setex(grant_key, self.GRANT_TTL, aes_key_hex)

        return aes_key_hex

    def list_reports_by_author(self, author_address: str) -> list[dict]:
        """List all report tx hashes and metadata for an author."""
        index_key = f"{self.INDEX_PREFIX}{author_address.lower()}"
        tx_hashes = self.redis.zrevrange(index_key, 0, -1)

        reports = []
        for tx_hash in tx_hashes:
            if isinstance(tx_hash, bytes):
                tx_hash = tx_hash.decode()
            meta = self._get_metadata(tx_hash)
            if meta:
                # Don't expose the AES key in listings
                safe_meta = {k: v for k, v in meta.items() if k != "aes_key_hex"}
                reports.append(safe_meta)

        return reports

    def _get_metadata(self, tx_hash: str) -> Optional[dict]:
        """Get report metadata from Redis."""
        raw = self.redis.get(f"{self.KEY_PREFIX}{tx_hash}")
        if not raw:
            return None
        if isinstance(raw, bytes):
            raw = raw.decode()
        return json.loads(raw)

    def _is_pledged(self, address: str) -> bool:
        """Check if an address is a pledged syndicate member."""
        try:
            with open(REGISTRY_PATH, "r") as f:
                registry = yaml.safe_load(f)

            agents = registry.get("agents") or []
            for agent in agents:
                if agent.get("base_address", "").lower() == address.lower():
                    if agent.get("status") == "active":
                        return True
            return False
        except Exception:
            return False