"""
MeshBay Node — content replication (node-to-node, admin-authorized).

A replication node downloads files from a source node and stores them
locally, then registers itself as an additional swarm source in the hub.
This provides redundancy and improves availability for public content.

Only public content is replicated (no GEK needed).
Private content replication requires the GEK and is admin-controlled.

Usage:
    replicator = ContentReplicator(
        hub_url=..., access_token=..., source_endpoint=...,
        local_dir=Path("/data/replicated"), node_pk_b64=...,
    )
    await replicator.replicate_file(file_id, file_name, file_size)
"""

import asyncio
import hashlib
import logging
from pathlib import Path

import blake3
import httpx

log = logging.getLogger(__name__)

CHUNK_SIZE = 1024 * 1024   # 1 MB


class ContentReplicator:
    """
    Downloads public files from a source node and registers as swarm source.
    """

    def __init__(
        self,
        hub_url: str,
        access_token: str,
        source_endpoint: str,    # "http://ip:port" of source node HTTP API
        local_dir: Path,
        node_pk_b64: str,
    ):
        self._hub_url        = hub_url.rstrip("/")
        self._access_token   = access_token
        self._source_endpoint = source_endpoint.rstrip("/")
        self._local_dir      = local_dir
        self._node_pk_b64    = node_pk_b64
        self._local_dir.mkdir(parents=True, exist_ok=True)

    @property
    def _auth_headers(self) -> dict:
        return {"Authorization": f"Bearer {self._access_token}"}

    async def fetch_index(self) -> list[dict]:
        """Fetch the public Mesh Group Index from the source node."""
        async with httpx.AsyncClient(timeout=30) as c:
            r = await c.get(f"{self._source_endpoint}/index")
            r.raise_for_status()
            return r.json()["entries"]

    async def replicate_file(
        self,
        file_id: str,
        file_name: str,
        file_size: int,
        progress_cb=None,
    ) -> Path:
        """
        Download a public file from the source node, verify integrity,
        save locally, and register as swarm source in the hub.
        Returns the local file path.
        """
        local_path = self._local_dir / file_name
        if local_path.exists():
            # Verify hash
            existing_hash = blake3.blake3(local_path.read_bytes()).hexdigest()
            if existing_hash == file_id:
                log.info("Already have %s, skipping", file_name)
                await self._register_swarm(file_id, local_path)
                return local_path

        log.info("Replicating %s (%d bytes) from %s", file_name, file_size, self._source_endpoint)

        # Stream download chunk by chunk
        n_chunks = max(1, (file_size + CHUNK_SIZE - 1) // CHUNK_SIZE)
        with open(local_path, "wb") as f:
            async with httpx.AsyncClient(timeout=60) as c:
                for chunk_idx in range(n_chunks):
                    # Download full file (simpler for public content)
                    if chunk_idx == 0:
                        r = await c.get(
                            f"{self._source_endpoint}/file/{file_id}",
                            headers=self._auth_headers,
                        )
                        r.raise_for_status()
                        f.write(r.content)
                        if progress_cb:
                            progress_cb(len(r.content), file_size)
                        break   # full file downloaded in one request

        # Verify hash
        actual_hash = blake3.blake3(local_path.read_bytes()).hexdigest()
        if actual_hash != file_id:
            local_path.unlink(missing_ok=True)
            raise ValueError(f"Hash mismatch: expected {file_id[:16]}, got {actual_hash[:16]}")

        log.info("Replicated %s (hash OK)", file_name)
        await self._register_swarm(file_id, local_path)
        return local_path

    async def _register_swarm(self, content_hash: str, local_path: Path) -> None:
        """Register this node as a swarm source for the content hash in the hub."""
        try:
            async with httpx.AsyncClient(timeout=10) as c:
                r = await c.post(
                    f"{self._hub_url}/v1/swarm/register",
                    json={"content_hash": content_hash, "endpoint": self._source_endpoint},
                    headers=self._auth_headers,
                )
                r.raise_for_status()
                log.debug("Registered as swarm source for %s", content_hash[:16])
        except Exception as e:
            log.warning("Failed to register swarm source: %s", e)

    async def replicate_all(self, progress_cb=None) -> list[Path]:
        """Replicate all public files from the source node."""
        entries = await self.fetch_index()
        results = []
        for entry in entries:
            try:
                path = await self.replicate_file(
                    file_id=entry["id"],
                    file_name=entry["name"],
                    file_size=entry["size"],
                    progress_cb=progress_cb,
                )
                results.append(path)
            except Exception as e:
                log.error("Failed to replicate %s: %s", entry["name"], e)
        return results