"""Federation — node identity, signed feeds, peer registry.

Identity layer for internet-wide federation of psyc nodes. Each node owns
an Ed25519 keypair persisted under DATA_DIR/federation/. The public key
fingerprint (first 16 bytes of SHA256(raw_pubkey) hex-encoded) goes into a
DNS TXT record so peers can discover and authenticate the node, and the
private key signs the outbound feed at /federation/feed.

This module is the *identity* primitives only — discovery walkers,
vouching/quorum, transparency log and auto-pull live in later stages.
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional, Tuple

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519
from pydantic import BaseModel, Field

from psyc import DATA_DIR, db, log
from psyc.result import Err, Ok, Result


_log = log.get(__name__)

FED_DIR = DATA_DIR / "federation"
PRIVATE_KEY_PATH = FED_DIR / "node.key"
PUBLIC_KEY_PATH = FED_DIR / "node.pub"

FEED_VERSION = "psyc1"
FEED_ALG = "ed25519"
FEED_PATH = "/federation/feed"


# ---------- keypair persistence -----------------------------------------

def _ensure_dir() -> None:
    FED_DIR.mkdir(parents=True, exist_ok=True)


def node_keypair() -> Tuple[ed25519.Ed25519PrivateKey, ed25519.Ed25519PublicKey]:
    """Return the node's Ed25519 keypair, generating + persisting it on first call.

    Private key lands at data/federation/node.key (PEM, chmod 0600); public
    at data/federation/node.pub (PEM). Idempotent — subsequent calls load
    the existing files instead of generating new ones.
    """
    _ensure_dir()
    if PRIVATE_KEY_PATH.exists() and PUBLIC_KEY_PATH.exists():
        priv_pem = PRIVATE_KEY_PATH.read_bytes()
        priv = serialization.load_pem_private_key(priv_pem, password=None)
        if not isinstance(priv, ed25519.Ed25519PrivateKey):
            raise RuntimeError(f"federation key at {PRIVATE_KEY_PATH} is not Ed25519")
        return priv, priv.public_key()

    priv = ed25519.Ed25519PrivateKey.generate()
    priv_pem = priv.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )
    pub = priv.public_key()
    pub_pem = pub.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )
    PRIVATE_KEY_PATH.write_bytes(priv_pem)
    os.chmod(PRIVATE_KEY_PATH, 0o600)
    PUBLIC_KEY_PATH.write_bytes(pub_pem)
    _log.info("federation.keypair.generated", path=str(PRIVATE_KEY_PATH))
    return priv, pub


def public_key_pem() -> str:
    """PEM-encoded public key as text — what peers store + verify against."""
    _, pub = node_keypair()
    return pub.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    ).decode("ascii")


def _raw_pubkey_bytes(pub: ed25519.Ed25519PublicKey) -> bytes:
    return pub.public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )


def node_fingerprint() -> str:
    """Short stable id for the node — first 16 bytes of SHA256(raw_pubkey), hex.

    Lives in DNS TXT records; 32 hex chars is short enough to fit but long
    enough to be collision-safe for any plausible peer population.
    """
    _, pub = node_keypair()
    digest = hashlib.sha256(_raw_pubkey_bytes(pub)).digest()
    return digest[:16].hex()


def _fingerprint_for_pubkey_pem(pubkey_pem: str) -> str:
    pub = serialization.load_pem_public_key(pubkey_pem.encode("ascii"))
    if not isinstance(pub, ed25519.Ed25519PublicKey):
        raise ValueError("not an Ed25519 public key")
    return hashlib.sha256(_raw_pubkey_bytes(pub)).digest()[:16].hex()


# ---------- DNS record format -------------------------------------------

class DNSRecord(BaseModel):
    """The SRV + TXT pair an admin pastes into their zone file."""
    srv_name: str
    srv_target: str
    srv_port: int
    srv_priority: int = 10
    srv_weight: int = 10
    txt_name: str
    txt_value: str
    human_instructions: str


def dns_record(domain: str, port: int = 443) -> DNSRecord:
    """Build the DNS-SD-style records that advertise this node at `domain`."""
    fp = node_fingerprint()
    srv_name = f"_psyc._tcp.{domain}"
    srv_target = f"{domain}."
    txt_name = srv_name
    txt_value = f"v={FEED_VERSION} fp={fp} alg={FEED_ALG} path={FEED_PATH}"
    instructions = (
        f"; psyc federation records for {domain}\n"
        f"; ----------------------------------------------------------\n"
        f"; 1) SRV record — locates this psyc node (host + port).\n"
        f'{srv_name}.  3600  IN  SRV  10 10 {port} {srv_target}\n'
        f";\n"
        f"; 2) TXT record — declares protocol version, key fingerprint,\n"
        f";    signature algorithm, and the feed endpoint path.\n"
        f'{txt_name}.  3600  IN  TXT  "{txt_value}"\n'
        f"; ----------------------------------------------------------\n"
        f"; Once these are live, federation peers can fetch:\n"
        f";   https://{domain}{FEED_PATH}   (signed feed JSON)\n"
        f";   https://{domain}/federation/key   (public key PEM)\n"
        f";   https://{domain}/federation/info  (capabilities)\n"
    )
    return DNSRecord(
        srv_name=srv_name,
        srv_target=srv_target,
        srv_port=port,
        txt_name=txt_name,
        txt_value=txt_value,
        human_instructions=instructions,
    )


# ---------- signing -----------------------------------------------------

def canonical_json(obj: Dict[str, Any]) -> bytes:
    """Deterministic JSON serialization — what we sign + hash over."""
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def sign_payload(payload: bytes) -> bytes:
    """Ed25519 signature over `payload`. Raw 64-byte sig."""
    priv, _ = node_keypair()
    return priv.sign(payload)


def verify_payload(payload: bytes, signature: bytes, pubkey_pem: str) -> bool:
    """True iff `signature` verifies under `pubkey_pem`. Never raises."""
    try:
        pub = serialization.load_pem_public_key(pubkey_pem.encode("ascii"))
        if not isinstance(pub, ed25519.Ed25519PublicKey):
            return False
        pub.verify(signature, payload)
        return True
    except Exception:
        return False


# ---------- feed export -------------------------------------------------

def _case_digest(case_record: Dict[str, Any]) -> str:
    return hashlib.sha256(canonical_json(case_record)).hexdigest()


def _build_case_records(window_hours: int) -> List[Dict[str, Any]]:
    cutoff = datetime.now(timezone.utc) - timedelta(hours=window_hours)
    out: List[Dict[str, Any]] = []
    for case in db.list_cases(limit=10_000):
        if case.ingested_at < cutoff:
            continue
        record: Dict[str, Any] = {
            "case_id": case.case_id,
            "summary": case.summary,
            "severity": case.classification.severity.value if case.classification.severity else None,
            "incident_type": case.classification.incident_type.value if case.classification.incident_type else None,
            "observed_at": case.observed_at.isoformat(),
            "feed_source": case.source_metadata.get("feed", ""),
            "iocs": (
                [{"value": v, "type": "url"} for v in case.observables.urls]
                + [{"value": v, "type": "domain"} for v in case.observables.domains]
                + [{"value": v, "type": "ip"} for v in case.observables.ips]
                + [{"value": v, "type": "hash"} for v in case.observables.hashes]
                + [{"value": v, "type": "cve"} for v in case.observables.cves]
            ),
        }
        record["digest_sha256"] = _case_digest(
            {k: v for k, v in record.items() if k != "digest_sha256"}
        )
        out.append(record)
    return out


def _build_ioc_records(window_hours: int) -> List[Dict[str, Any]]:
    cutoff = datetime.now(timezone.utc) - timedelta(hours=window_hours)
    out: List[Dict[str, Any]] = []
    seen: set = set()
    for ioc_type in ("url", "domain", "ip", "hash", "cve"):
        for row in db.iocs_by_type(ioc_type):
            first_seen = row.get("first_seen")
            if first_seen:
                try:
                    if datetime.fromisoformat(first_seen) < cutoff:
                        continue
                except ValueError:
                    pass
            key = (row["value"], row["ioc_type"])
            if key in seen:
                continue
            seen.add(key)
            record = {
                "value": row["value"],
                "type": row["ioc_type"],
                "severity": row.get("severity"),
                "first_seen": first_seen,
            }
            record["digest_sha256"] = hashlib.sha256(canonical_json(record)).hexdigest()
            out.append(record)
    return out


def build_signed_feed(window_hours: int = 24) -> Dict[str, Any]:
    """Build the JSON feed peers will pull from /federation/feed.

    Pulls cases ingested in the last `window_hours` plus the corresponding
    IOC slice, attaches per-record `digest_sha256` (so peers can later
    quorum-match across nodes), and signs the canonical JSON of the whole
    payload-minus-signature with our Ed25519 key.
    """
    payload: Dict[str, Any] = {
        "version": FEED_VERSION,
        "fingerprint": node_fingerprint(),
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "window_hours": window_hours,
        "cases": _build_case_records(window_hours),
        "iocs": _build_ioc_records(window_hours),
    }
    sig = sign_payload(canonical_json(payload))
    payload["signature"] = base64.b64encode(sig).decode("ascii")
    return payload


# ---------- import + quorum-signal buffer -------------------------------

class ImportSummary(BaseModel):
    peer_fingerprint: str
    cases_seen: int
    iocs_seen: int
    signal_ids: List[Tuple[str, str]] = Field(default_factory=list)


def import_signed_feed(feed: Dict[str, Any], expected_pubkey_pem: str) -> Result[ImportSummary, str]:
    """Verify + record a peer's feed into the federation_signals buffer.

    Does NOT merge into the local case store — that's the quorum stage's
    job. The buffer is the per-hash signal log that quorum logic later
    aggregates ("3 trusted peers reported this same IOC → promote").
    """
    sig_b64 = feed.get("signature")
    if not sig_b64:
        return Err("missing signature")
    try:
        signature = base64.b64decode(sig_b64)
    except Exception:
        return Err("malformed signature (not base64)")

    unsigned = {k: v for k, v in feed.items() if k != "signature"}
    if not verify_payload(canonical_json(unsigned), signature, expected_pubkey_pem):
        return Err("signature verification failed")

    peer_fp = feed.get("fingerprint", "")
    if not peer_fp:
        return Err("missing fingerprint")
    if peer_fp == node_fingerprint():
        return Err("loop: own feed")

    # Cross-check the declared fingerprint matches the pubkey we verified with.
    try:
        if _fingerprint_for_pubkey_pem(expected_pubkey_pem) != peer_fp:
            return Err("fingerprint does not match provided pubkey")
    except Exception as exc:
        return Err(f"bad pubkey: {exc}")

    now = datetime.now(timezone.utc).isoformat()
    signal_ids: List[Tuple[str, str]] = []
    cases = feed.get("cases") or []
    iocs = feed.get("iocs") or []

    for c in cases:
        case_id = c.get("case_id") or ""
        digest = c.get("digest_sha256") or hashlib.sha256(canonical_json(c)).hexdigest()
        db.record_signal(dict(
            peer_fingerprint=peer_fp,
            signal_type="case",
            signal_id=case_id,
            signal_hash=digest,
            received_at=now,
            raw_json=json.dumps(c, sort_keys=True),
        ))
        signal_ids.append(("case", digest))

    for i in iocs:
        value = i.get("value") or ""
        digest = i.get("digest_sha256") or hashlib.sha256(canonical_json(i)).hexdigest()
        db.record_signal(dict(
            peer_fingerprint=peer_fp,
            signal_type="ioc",
            signal_id=value,
            signal_hash=digest,
            received_at=now,
            raw_json=json.dumps(i, sort_keys=True),
        ))
        signal_ids.append(("ioc", digest))

    _log.info("federation.import.ok", peer=peer_fp, cases=len(cases), iocs=len(iocs))
    return Ok(ImportSummary(
        peer_fingerprint=peer_fp,
        cases_seen=len(cases),
        iocs_seen=len(iocs),
        signal_ids=signal_ids,
    ))


# ---------- peer registry ------------------------------------------------

class Peer(BaseModel):
    domain: str
    fingerprint: str
    pubkey_pem: str
    status: str = "unknown"  # unknown | trusted | blocked
    discovered_at: str
    last_seen: Optional[str] = None
    notes: Optional[str] = None


def _row_to_peer(row: Dict[str, Any]) -> Peer:
    return Peer(
        domain=row["domain"],
        fingerprint=row["fingerprint"],
        pubkey_pem=row["pubkey_pem"],
        status=row.get("status") or "unknown",
        discovered_at=row.get("discovered_at") or "",
        last_seen=row.get("last_seen"),
        notes=row.get("notes"),
    )


def register_peer(domain: str, fingerprint: str, pubkey_pem: str, status: str = "unknown") -> None:
    """Insert or update a peer in the registry. Idempotent on `domain`."""
    now = datetime.now(timezone.utc).isoformat()
    existing = db.get_peer(domain)
    discovered_at = existing["discovered_at"] if existing else now
    db.upsert_peer(dict(
        domain=domain,
        fingerprint=fingerprint,
        pubkey_pem=pubkey_pem,
        status=status,
        discovered_at=discovered_at,
        last_seen=now,
        notes=existing.get("notes") if existing else None,
    ))


def list_peers() -> List[Peer]:
    return [_row_to_peer(r) for r in db.list_peers()]


def get_peer(domain: str) -> Optional[Peer]:
    row = db.get_peer(domain)
    return _row_to_peer(row) if row else None


def set_peer_status(domain: str, status: str) -> None:
    if status not in ("unknown", "trusted", "blocked"):
        raise ValueError(f"unknown peer status: {status}")
    db.set_peer_status(domain, status)


def remove_peer(domain: str) -> None:
    db.remove_peer(domain)
