"""PeerDrop Core — File transfer via stream handshake + bitswap."""

from __future__ import annotations

import json
import logging
import os
import uuid
from datetime import datetime
from typing import TYPE_CHECKING, Any

import trio

from peerdrop.core.events import (
    EventBus,
    TransferCompleted,
    TransferFailed,
    TransferProgress,
    TransferStarted,
)
from peerdrop.core.models import Transfer, TransferRequest, TransferStatus

if TYPE_CHECKING:
    from py_ipfs_lite.peer import Peer as IPFSPeer

logger = logging.getLogger(__name__)

TRANSFER_PROTOCOL = "/peerdrop/transfer/1.0.0"

# Per-block fetch timeout (seconds) — generous for slow networks
BLOCK_TIMEOUT = 120.0

# Overall transfer timeout (seconds) — 30 minutes max for any file
OVERALL_TRANSFER_TIMEOUT = 1800.0

# How long the sender waits for the receiver to ACK the handshake
ACK_TIMEOUT = 60.0


def _serialize_request(req: TransferRequest) -> bytes:
    """Serialize a TransferRequest to JSON bytes."""
    return json.dumps({
        "file_name": req.file_name,
        "file_size": req.file_size,
        "root_cid": req.root_cid,
        "sender_peer_id": req.sender_peer_id,
    }).encode()


def _deserialize_request(data: bytes) -> TransferRequest:
    """Deserialize JSON bytes to a TransferRequest."""
    obj = json.loads(data)
    return TransferRequest(
        file_name=obj["file_name"],
        file_size=obj["file_size"],
        root_cid=obj["root_cid"],
        sender_peer_id=obj["sender_peer_id"],
    )


class TransferManager:
    """Manages file transfers between peers.

    Transfer flow:
    1. Sender calls send_file() which:
       - Adds file to local blockstore via py-ipfs-lite (chunking + CID)
       - Opens stream to receiver, sends TransferRequest
       - Waits for ACK
    2. Receiver's stream handler:
       - Reads TransferRequest
       - Sends ACK
       - Fetches blocks via py-ipfs-lite get_file(root_cid)
       - Writes file to disk
       - Sends DONE confirmation back on same stream
    3. Sender receives DONE, marks transfer as completed.
    """

    def __init__(self, ipfs_peer: IPFSPeer, event_bus: EventBus, download_dir: str = "~/.peerdrop/downloads") -> None:
        self._ipfs_peer = ipfs_peer
        self._event_bus = event_bus
        self._transfers: dict[str, Transfer] = {}
        self._cancel_scopes: dict[str, trio.CancelScope] = {}
        self._download_dir = os.path.expanduser(download_dir)
        os.makedirs(self._download_dir, exist_ok=True)

    @property
    def download_dir(self) -> str:
        """Return the current download directory."""
        return self._download_dir

    def set_download_dir(self, new_dir: str) -> str:
        """Change the download directory. Returns the resolved path."""
        self._download_dir = os.path.expanduser(new_dir)
        os.makedirs(self._download_dir, exist_ok=True)
        return self._download_dir

    async def send_file(self, file_path: str, target_peer_id: str) -> Transfer:
        """Send a file to a target peer.

        Steps:
        1. Add file to local blockstore (py-ipfs-lite handles chunking)
        2. Open stream to receiver
        3. Send TransferRequest with root CID
        4. Wait for ACK, then for progress updates and a DONE confirmation.

        Only marks the transfer COMPLETED when the receiver explicitly confirms
        with ``DONE``. Any failure or timeout along the way marks it FAILED.
        """
        from libp2p.peer.id import ID

        transfer = Transfer(
            transfer_id=uuid.uuid4().hex[:12],
            file_path=file_path,
            file_name=os.path.basename(file_path),
            sender_peer_id=str(self._ipfs_peer.host.id()),
            receiver_peer_id=target_peer_id,
            status=TransferStatus.CONNECTING,
        )
        self._transfers[transfer.transfer_id] = transfer

        try:
            # Step 1: Add file to local blockstore
            logger.info(f"Adding file to blockstore: {file_path}")
            node = await self._ipfs_peer.add_file(file_path)
            root_cid = str(node.cid())
            transfer.root_cid = root_cid
            # py-ipfs-lite's add_file() returns an IPLDNode holding only the CID,
            # so node.raw_data() is always None. Take the real size from the file.
            try:
                transfer.file_size = os.path.getsize(file_path)
            except OSError:
                transfer.file_size = 0

            # Step 2: Open stream to receiver
            target_id = ID.from_string(target_peer_id)
            logger.info(f"Opening stream to {target_peer_id}")
            stream = await self._ipfs_peer.host.open_stream(
                target_id, [TRANSFER_PROTOCOL]
            )

            # Step 3: Send TransferRequest
            request = TransferRequest(
                file_name=transfer.file_name,
                file_size=transfer.file_size,
                root_cid=root_cid,
                sender_peer_id=transfer.sender_peer_id,
            )
            await stream.write(_serialize_request(request))

            # Step 4: Wait for ACK/REJECT (bounded so we can't hang forever)
            with trio.fail_after(ACK_TIMEOUT):
                response = (await stream.read()).strip()

            if response == b"REJECT":
                await stream.close()
                transfer.status = TransferStatus.FAILED
                transfer.error = "Receiver rejected transfer"
                self._event_bus.publish(TransferFailed(
                    transfer_id=transfer.transfer_id,
                    error="Receiver rejected transfer",
                ))
                return transfer

            if response != b"ACK":
                await stream.close()
                transfer.status = TransferStatus.FAILED
                transfer.error = f"Unexpected receiver response: {response!r}"
                self._event_bus.publish(TransferFailed(
                    transfer_id=transfer.transfer_id,
                    error=transfer.error,
                ))
                return transfer

            # The user may have cancelled while the handshake was in flight —
            # don't overwrite the CANCELLED status.
            if transfer.status == TransferStatus.CANCELLED:
                logger.info(f"Transfer cancelled during handshake: {transfer.transfer_id}")
                await stream.close()
                return transfer

            # Transfer accepted - wait for progress updates and DONE
            transfer.status = TransferStatus.TRANSFERRING
            transfer.started_at = datetime.now()
            self._event_bus.publish(TransferStarted(
                transfer_id=transfer.transfer_id,
                file_name=transfer.file_name,
                sender_peer_id=transfer.sender_peer_id,
                receiver_peer_id=target_peer_id,
            ))

            logger.info("Waiting for receiver to confirm file delivery...")

            # Confirmation messages are newline-delimited (PROGRESS:<float>\n,
            # DONE\n) so that partial or concatenated frames are parsed safely.
            read_buffer = b""

            def _next_message() -> bytes | None:
                nonlocal read_buffer
                idx = read_buffer.find(b"\n")
                if idx == -1:
                    return None
                line, read_buffer = read_buffer[:idx], read_buffer[idx + 1:]
                return line

            completed = False
            try:
                with trio.fail_after(OVERALL_TRANSFER_TIMEOUT):
                    while True:
                        # Respect user cancellation even while blocked on read()
                        if transfer.status == TransferStatus.CANCELLED:
                            logger.info(f"Transfer cancelled while awaiting confirmation: {transfer.transfer_id}")
                            break

                        with trio.move_on_after(1.0):
                            chunk = await stream.read()
                            read_buffer += chunk

                        while True:
                            confirmation = _next_message()
                            if confirmation is None:
                                break
                            if confirmation == b"DONE":
                                logger.info("Receiver confirmed file delivery")
                                completed = True
                                break
                            elif confirmation.startswith(b"PROGRESS:"):
                                # Receiver sends progress updates: PROGRESS:<progress_float>
                                try:
                                    progress_val = float(confirmation.split(b":", 1)[1])
                                    transfer.progress = progress_val
                                    self._event_bus.publish(TransferProgress(
                                        transfer_id=transfer.transfer_id,
                                        progress=progress_val,
                                        bytes_sent=int(progress_val * transfer.file_size),
                                    ))
                                except (ValueError, IndexError):
                                    pass
                            else:
                                logger.warning(f"Unexpected confirmation: {confirmation}")
                        if completed:
                            break
            except trio.TooSlowError:
                transfer.error = "Timed out waiting for receiver confirmation"
                logger.error(f"Transfer timed out: {transfer.transfer_id}: {transfer.error}")
            except Exception as e:
                transfer.error = str(e)
                logger.error(f"Transfer failed while awaiting confirmation: {transfer.transfer_id}: {e}")
            finally:
                await stream.close()

            if transfer.status == TransferStatus.CANCELLED:
                return transfer

            if completed:
                transfer.status = TransferStatus.COMPLETED
                transfer.progress = 1.0
                transfer.completed_at = datetime.now()
                self._event_bus.publish(TransferCompleted(
                    transfer_id=transfer.transfer_id,
                    file_name=transfer.file_name,
                ))
                logger.info(f"Transfer completed: {transfer.transfer_id}")
            else:
                transfer.status = TransferStatus.FAILED
                transfer.error = transfer.error or "Transfer failed before receiver confirmed"
                self._event_bus.publish(TransferFailed(
                    transfer_id=transfer.transfer_id,
                    error=transfer.error,
                ))
                logger.error(f"Transfer failed: {transfer.transfer_id}: {transfer.error}")

            return transfer

        except Exception as e:
            transfer.status = TransferStatus.FAILED
            transfer.error = str(e)
            self._event_bus.publish(TransferFailed(
                transfer_id=transfer.transfer_id,
                error=str(e),
            ))
            logger.error(f"Transfer failed: {e}")
            return transfer

    async def handle_incoming_transfer(self, stream: Any) -> None:
        """Handle an incoming transfer request from a stream.

        This is registered as a stream handler on TRANSFER_PROTOCOL.
        """
        transfer: Transfer | None = None
        output_path: str | None = None

        try:
            # Read TransferRequest
            data = await stream.read()
            request = _deserialize_request(data)

            logger.info(
                f"Incoming transfer: {request.file_name} "
                f"({request.file_size} bytes) from {request.sender_peer_id}"
            )

            # Sanitize the filename so a malicious sender cannot escape the
            # download directory via path traversal.
            safe_name = os.path.basename(request.file_name) or "unknown"

            transfer = Transfer(
                transfer_id=uuid.uuid4().hex[:12],
                file_name=safe_name,
                file_size=request.file_size,
                root_cid=request.root_cid,
                sender_peer_id=request.sender_peer_id,
                receiver_peer_id=str(self._ipfs_peer.host.id()),
                status=TransferStatus.TRANSFERRING,
                started_at=datetime.now(),
            )
            self._transfers[transfer.transfer_id] = transfer

            # Send ACK (newline-delimited so the sender can parse safely)
            await stream.write(b"ACK\n")

            self._event_bus.publish(TransferStarted(
                transfer_id=transfer.transfer_id,
                file_name=safe_name,
                sender_peer_id=request.sender_peer_id,
                receiver_peer_id=transfer.receiver_peer_id,
            ))

            # Fetch file via py-ipfs-lite (bitswap handles the block transfer)
            # Use output_path= for streaming to disk — avoids buffering the
            # entire file in memory.
            logger.info(f"Fetching file via bitswap: {request.root_cid}")
            from libp2p.peer.id import ID as PeerID
            sender_id = PeerID.from_base58(request.sender_peer_id)
            host = self._ipfs_peer.host
            sender_addrs = host.get_peerstore().addrs(sender_id)
            provider_addr = None
            if sender_addrs:
                provider_addr = str(sender_addrs[0])
                logger.info(f"Using sender address for bitswap: {provider_addr}")

            output_path = os.path.join(self._download_dir, safe_name)

            def _progress_callback(bytes_done: int, total: int) -> None:
                # py-ipfs-lite only calls this when it can derive the total file
                # size from its LOCAL blockstore — which the receiver doesn't
                # have yet during a fresh remote fetch, so this is best-effort.
                progress_val = bytes_done / total if total > 0 else 0.0
                transfer.progress = progress_val
                transfer.bytes_sent = bytes_done
                self._event_bus.publish(TransferProgress(
                    transfer_id=transfer.transfer_id,
                    progress=progress_val,
                    bytes_sent=bytes_done,
                ))

            # Register a cancel scope so cancel_transfer() can abort the fetch.
            cancel_scope = trio.CancelScope()
            self._cancel_scopes[transfer.transfer_id] = cancel_scope
            try:
                with cancel_scope:
                    with trio.fail_after(OVERALL_TRANSFER_TIMEOUT):
                        async with trio.open_nursery() as progress_nursery:
                            # Background task: periodically report progress to the
                            # sender. Progress is derived from bytes written to
                            # disk (the sender's reported file_size is accurate),
                            # not from py-ipfs-lite's progress_callback, which is
                            # never invoked for fresh remote fetches.
                            async def _report_progress() -> None:
                                last_sent = -1.0
                                total = transfer.file_size or 1
                                while True:
                                    try:
                                        written = os.path.getsize(output_path)
                                    except OSError:
                                        written = 0
                                    p = min(1.0, written / total)
                                    transfer.progress = p
                                    transfer.bytes_sent = written
                                    # Throttle: send every 5% change or at 100%
                                    if p - last_sent >= 0.05 or p >= 1.0:
                                        last_sent = p
                                        try:
                                            await stream.write(
                                                b"PROGRESS:" + str(p).encode() + b"\n"
                                            )
                                        except Exception:
                                            break
                                    await trio.sleep(0.25)

                            progress_nursery.start_soon(_report_progress)
                            try:
                                await self._ipfs_peer.get_file(
                                    request.root_cid,
                                    output_path=output_path,
                                    provider_addr=provider_addr,
                                    timeout=BLOCK_TIMEOUT,
                                    progress_callback=_progress_callback,
                                )
                            finally:
                                progress_nursery.cancel_scope.cancel()
            finally:
                self._cancel_scopes.pop(transfer.transfer_id, None)

            transfer.status = TransferStatus.COMPLETED
            transfer.progress = 1.0
            transfer.completed_at = datetime.now()

            # Send DONE confirmation back to sender
            try:
                await stream.write(b"DONE\n")
                logger.info("Sent DONE confirmation to sender")
            except Exception as e:
                logger.warning(f"Failed to send DONE confirmation: {e}")

            self._event_bus.publish(TransferCompleted(
                transfer_id=transfer.transfer_id,
                file_name=safe_name,
            ))

            logger.info(f"Transfer received: {output_path}")

        except trio.Cancelled:
            # Cancelled via cancel_transfer() — our cancel scope was triggered.
            if transfer is not None:
                if transfer.status != TransferStatus.CANCELLED:
                    transfer.status = TransferStatus.CANCELLED
                transfer.error = "Transfer cancelled"
                logger.info(f"Incoming transfer cancelled: {transfer.transfer_id}")
            try:
                await stream.reset()
            except Exception:
                pass
            self._remove_partial_output(output_path)

        except Exception as e:
            logger.error(f"Incoming transfer failed: {e}")
            if transfer is not None:
                transfer.status = TransferStatus.FAILED
                transfer.error = str(e)
            try:
                await stream.reset()
            except Exception:
                pass
            self._remove_partial_output(output_path)
            self._event_bus.publish(TransferFailed(
                transfer_id=transfer.transfer_id if transfer else "unknown",
                error=str(e),
            ))

    @staticmethod
    def _remove_partial_output(output_path: str | None) -> None:
        """Remove a partially-written download on failure/cancellation."""
        if output_path:
            try:
                if os.path.exists(output_path):
                    os.remove(output_path)
            except OSError:
                pass

    def get_transfer(self, transfer_id: str) -> Transfer | None:
        """Lookup a transfer by ID."""
        return self._transfers.get(transfer_id)

    def list_transfers(self) -> list[Transfer]:
        """Return all transfers."""
        return list(self._transfers.values())

    def cancel_transfer(self, transfer_id: str) -> bool:
        """Cancel an active transfer. Returns True if cancelled.

        Sets the status AND aborts in-flight work: for incoming transfers this
        cancels the bitswap fetch scope; for outgoing transfers the sender's
        wait loop notices the status change and stops.
        """
        transfer = self._transfers.get(transfer_id)
        if not transfer:
            return False

        if transfer.status in (TransferStatus.COMPLETED, TransferStatus.FAILED):
            return False

        transfer.status = TransferStatus.CANCELLED
        transfer.completed_at = datetime.now()

        # Abort in-flight fetch (receiver side)
        scope = self._cancel_scopes.pop(transfer_id, None)
        if scope is not None:
            scope.cancel()

        logger.info(f"Transfer cancelled: {transfer_id}")
        return True