"""File download handler: DMS → local filesystem.

Receives chunked file transfers from qRunX over QRP:
  download_file_start   → create target file, prepare buffer
  download_file_chunk   → append base64-decoded data
  download_file_complete → verify checksum, finalise

Architecture reference: Section 11.3 (File Transfer Flow).
"""

from __future__ import annotations

import base64
import hashlib
import logging
import os
import tempfile
from typing import Any

logger = logging.getLogger("qremotex.transfer.downloader")

# Maximum single file size to accept (512 MB)
MAX_FILE_SIZE_BYTES = 512 * 1024 * 1024


class _ActiveDownload:
    """Tracks state for a single in-progress file download."""

    __slots__ = (
        "file_id", "target_path", "tmp_path", "total_chunks",
        "total_size", "received_chunks", "received_bytes",
        "sha256", "fh",
    )

    def __init__(
        self,
        file_id: str,
        target_path: str,
        tmp_path: str,
        total_chunks: int,
        total_size: int,
    ) -> None:
        self.file_id = file_id
        self.target_path = target_path
        self.tmp_path = tmp_path
        self.total_chunks = total_chunks
        self.total_size = total_size
        self.received_chunks = 0
        self.received_bytes = 0
        self.sha256 = hashlib.sha256()
        self.fh: Any = None

    def close(self) -> None:
        if self.fh and not self.fh.closed:
            self.fh.close()

    def cleanup(self) -> None:
        """Remove temporary file on failure."""
        self.close()
        if os.path.exists(self.tmp_path):
            os.unlink(self.tmp_path)


class FileDownloader:
    """Manages incoming file downloads from qRunX/DMS."""

    def __init__(self, download_dir: str | None = None) -> None:
        self._download_dir = download_dir or tempfile.gettempdir()
        os.makedirs(self._download_dir, exist_ok=True)
        # Active downloads keyed by file_id
        self._active: dict[str, _ActiveDownload] = {}

    @property
    def active_count(self) -> int:
        return len(self._active)

    def handle_start(self, msg: dict[str, Any]) -> dict[str, Any]:
        """Process download_file_start message.

        Returns:
            {"ok": True, "file_id": ...} on success,
            {"ok": False, "error": ...} on failure.
        """
        file_id = msg.get("file_id", "")
        target_path = msg.get("target_path", "")
        total_chunks = int(msg.get("total_chunks", 0))
        total_size = int(msg.get("total_size_bytes", 0))

        if not file_id or not target_path:
            return {"ok": False, "error": "Missing file_id or target_path"}

        if total_size > MAX_FILE_SIZE_BYTES:
            return {"ok": False, "error": f"File too large: {total_size} bytes (max {MAX_FILE_SIZE_BYTES})"}

        if file_id in self._active:
            # Duplicate start — cleanup previous attempt
            self._active[file_id].cleanup()

        # Validate target_path — prevent path traversal
        resolved = os.path.realpath(os.path.join(self._download_dir, target_path))
        if not resolved.startswith(os.path.realpath(self._download_dir)):
            return {"ok": False, "error": "Path traversal detected in target_path"}

        # Ensure parent directory exists
        parent = os.path.dirname(resolved)
        os.makedirs(parent, exist_ok=True)

        # Write to a temp file first, then move on completion
        tmp_path = resolved + ".qrx_download"

        try:
            fh = open(tmp_path, "wb")
        except OSError as e:
            return {"ok": False, "error": f"Cannot open file: {e}"}

        dl = _ActiveDownload(
            file_id=file_id,
            target_path=resolved,
            tmp_path=tmp_path,
            total_chunks=total_chunks,
            total_size=total_size,
        )
        dl.fh = fh
        self._active[file_id] = dl

        logger.info(
            "Download started: file_id=%s target=%s chunks=%d size=%d",
            file_id, target_path, total_chunks, total_size,
        )
        return {"ok": True, "file_id": file_id}

    def handle_chunk(self, msg: dict[str, Any]) -> dict[str, Any]:
        """Process download_file_chunk message."""
        file_id = msg.get("file_id", "")
        chunk_index = int(msg.get("chunk_index", -1))
        data_b64 = msg.get("data_b64", "")

        dl = self._active.get(file_id)
        if not dl:
            return {"ok": False, "error": f"No active download for file_id={file_id}"}

        if not data_b64:
            return {"ok": False, "error": "Empty chunk data"}

        try:
            raw = base64.b64decode(data_b64)
        except Exception as e:
            dl.cleanup()
            del self._active[file_id]
            return {"ok": False, "error": f"Invalid base64 data: {e}"}

        # Size guard
        if dl.received_bytes + len(raw) > MAX_FILE_SIZE_BYTES:
            dl.cleanup()
            del self._active[file_id]
            return {"ok": False, "error": "Download exceeds maximum file size"}

        try:
            dl.fh.write(raw)
            dl.sha256.update(raw)
            dl.received_bytes += len(raw)
            dl.received_chunks += 1
        except OSError as e:
            dl.cleanup()
            del self._active[file_id]
            return {"ok": False, "error": f"Write error: {e}"}

        return {"ok": True, "file_id": file_id, "chunk_index": chunk_index}

    def handle_complete(self, msg: dict[str, Any]) -> dict[str, Any]:
        """Process download_file_complete message.

        Verifies SHA-256 checksum and moves temp file to target.
        """
        file_id = msg.get("file_id", "")
        expected_checksum = msg.get("checksum", "")

        dl = self._active.get(file_id)
        if not dl:
            return {"ok": False, "error": f"No active download for file_id={file_id}"}

        dl.close()

        actual_checksum = f"sha256:{dl.sha256.hexdigest()}"

        if expected_checksum and actual_checksum != expected_checksum:
            dl.cleanup()
            del self._active[file_id]
            return {
                "ok": False,
                "error": f"Checksum mismatch: expected={expected_checksum} actual={actual_checksum}",
            }

        # Move temp file to final target
        try:
            os.replace(dl.tmp_path, dl.target_path)
        except OSError as e:
            dl.cleanup()
            del self._active[file_id]
            return {"ok": False, "error": f"Failed to finalize file: {e}"}

        del self._active[file_id]

        logger.info(
            "Download complete: file_id=%s path=%s bytes=%d chunks=%d checksum=%s",
            file_id, dl.target_path, dl.received_bytes, dl.received_chunks, actual_checksum,
        )
        return {
            "ok": True,
            "file_id": file_id,
            "local_path": dl.target_path,
            "size_bytes": dl.received_bytes,
            "checksum": actual_checksum,
        }

    def cleanup_all(self) -> None:
        """Abort all in-progress downloads."""
        for dl in self._active.values():
            dl.cleanup()
        self._active.clear()
