import aiofiles
import asyncio
import base64
import os
import struct
import time
import uuid
from pathlib import Path

from pantheon.toolset import tool
from pantheon.toolsets.file.file_manager import FileManagerToolSetBase
from pantheon.utils.log import logger


class FileTransferToolSet(FileManagerToolSetBase):
    """File transfer toolset with async I/O and smart flush strategy.

    This class provides basic file transfer functionality with performance optimizations:
    - Async I/O using aiofiles (non-blocking)
    - Smart flush strategy (balances performance and responsiveness)
    - Progress tracking for better UX

    Features:
    - open_file_for_write: Open file for async writing
    - write_chunk: Write data chunks with smart flushing
    - close_file: Close file with final flush
    - open_file_for_read: Open file for async reading
    - read_chunk: Read data chunks
    - read_file: Read entire file (streaming or non-streaming)
    """

    def __init__(
        self,
        name: str,
        path: str | Path,
        **kwargs,
    ):
        super().__init__(name, path, **kwargs)
        self._handles = {}  # {handle_id: aiofiles.File}
        self._file_info = {}  # {handle_id: metadata dict}

        # Smart flush strategy configuration
        self.FLUSH_INTERVAL_CHUNKS = 10  # Flush every N chunks
        self.FLUSH_INTERVAL_BYTES = 5 * 1024 * 1024  # Flush every 5MB
        self.FLUSH_INTERVAL_SECONDS = 2.0  # Flush every 2 seconds

    @tool
    async def open_file_for_write(self, file_path: str):
        """Open a file for writing (async).

        Args:
            file_path: Relative path to the file (cannot contain '..')

        Returns:
            dict: {"success": True, "handle_id": str} on success
                  {"success": False, "error": str} on failure
        """
        if ".." in file_path:
            return {"error": "File path cannot contain '..'"}

        path = self._get_root() / file_path

        # Ensure parent directory exists
        try:
            path.parent.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            logger.warning(f"Failed to create parent directory: {e}")

        handle_id = str(uuid.uuid4())

        try:
            # Open file asynchronously
            handle = await aiofiles.open(path, mode="wb")
            self._handles[handle_id] = handle
            self._file_info[handle_id] = {
                "path": path,
                "size": 0,
                "chunks_written": 0,
                "bytes_since_flush": 0,
                "last_flush_time": time.time(),
                "created_at": time.time(),
            }
            logger.debug(f"Opened file for write: {path} (handle_id={handle_id})")
            return {"success": True, "handle_id": handle_id}
        except Exception as e:
            logger.error(f"Failed to open file for write: {path}, error: {e}")
            return {"success": False, "error": str(e)}

    @tool
    async def write_chunk(self, handle_id: str, data):
        """Write a chunk to a file with smart flush strategy.

        Args:
            handle_id: File handle ID from open_file_for_write.
            data: Binary data as bytes (from cloudpickle) or
                  base64-encoded string (from JSON/frontend clients).

        Returns:
            dict: {"success": True, "bytes_written": int} on success
                  {"success": False, "error": str} on failure
        """
        if handle_id not in self._handles:
            return {"success": False, "error": "Handle not found"}

        handle = self._handles[handle_id]
        info = self._file_info[handle_id]

        # Decode data if base64-encoded string
        if isinstance(data, str):
            data = base64.b64decode(data)
        elif not isinstance(data, (bytes, bytearray)):
            return {
                "success": False,
                "error": f"Unsupported data type: {type(data).__name__}",
            }

        try:
            # Write data asynchronously (non-blocking)
            await handle.write(data)

            # Update statistics
            bytes_written = len(data)
            info["size"] += bytes_written
            info["chunks_written"] += 1
            info["bytes_since_flush"] += bytes_written

            # Smart flush strategy: flush if any condition is met
            current_time = time.time()
            last_flush_time = info["last_flush_time"]

            should_flush = (
                # Condition 1: Every N chunks
                info["chunks_written"] % self.FLUSH_INTERVAL_CHUNKS == 0
                or
                # Condition 2: Every N bytes
                info["bytes_since_flush"] >= self.FLUSH_INTERVAL_BYTES
                or
                # Condition 3: Every N seconds
                (current_time - last_flush_time) >= self.FLUSH_INTERVAL_SECONDS
            )

            if should_flush:
                await handle.flush()
                info["bytes_since_flush"] = 0
                info["last_flush_time"] = current_time
                logger.debug(
                    f"Flushed file (handle_id={handle_id}, "
                    f"chunks={info['chunks_written']}, "
                    f"total_size={info['size']} bytes)"
                )

            return {"success": True, "bytes_written": bytes_written}

        except Exception as e:
            logger.error(f"Failed to write chunk (handle_id={handle_id}): {e}")
            return {"success": False, "error": str(e)}

    @tool
    async def close_file(self, handle_id: str):
        """Close a file with final flush (async).

        Args:
            handle_id: File handle ID from open_file_for_write or open_file_for_read.

        Returns:
            dict: {"success": True, "total_size": int, "chunks_written": int} on success (write mode)
                  {"success": True} on success (read mode)
                  {"success": False, "error": str} on failure
        """
        if handle_id not in self._handles:
            return {"success": False, "error": "Handle not found"}

        handle = self._handles[handle_id]
        info = self._file_info[handle_id]

        try:
            # Final flush and close (async, non-blocking)
            await handle.flush()
            await handle.close()

            # Close the low-level fd used by read_chunk_at, if any.
            random_access_fd = info.get("random_access_fd")
            if random_access_fd is not None:
                try:
                    os.close(random_access_fd)
                except OSError as fd_err:
                    logger.warning(
                        f"Failed to close random-access fd for {info.get('path')}: {fd_err}"
                    )

            # Calculate statistics
            total_time = time.time() - info["created_at"]
            total_size = info["size"]

            # Check if this is a write handle (has chunks_written) or read handle
            is_write_handle = "chunks_written" in info

            if is_write_handle:
                chunks_written = info["chunks_written"]
                logger.info(
                    f"Closed file (write): {info['path']} "
                    f"(size={total_size} bytes, chunks={chunks_written}, time={total_time:.2f}s)"
                )
            else:
                logger.info(
                    f"Closed file (read): {info['path']} "
                    f"(size={total_size} bytes, time={total_time:.2f}s)"
                )

            # Cleanup
            del self._handles[handle_id]
            del self._file_info[handle_id]

            # Return appropriate response based on handle type
            if is_write_handle:
                return {
                    "success": True,
                    "total_size": total_size,
                    "chunks_written": chunks_written,
                    "duration_seconds": round(total_time, 2),
                }
            else:
                return {"success": True}

        except Exception as e:
            logger.error(f"Failed to close file (handle_id={handle_id}): {e}")

            # Cleanup on error
            try:
                await handle.close()
            except:
                pass

            if handle_id in self._handles:
                del self._handles[handle_id]
            if handle_id in self._file_info:
                del self._file_info[handle_id]

            return {"success": False, "error": str(e)}

    @tool
    async def open_file_for_read(self, file_path: str):
        """Open a file for chunked reading (async). Returns handle_id and total_size.

        Args:
            file_path: Relative path to the file (cannot contain '..')

        Returns:
            dict: {"success": True, "handle_id": str, "total_size": int} on success
                  {"success": False, "error": str} on failure
        """
        if ".." in file_path:
            return {"success": False, "error": "File path cannot contain '..'"}

        path = self._get_root() / file_path
        if not path.exists():
            return {"success": False, "error": "File does not exist"}

        handle_id = str(uuid.uuid4())

        try:
            # Open file asynchronously for reading
            handle = await aiofiles.open(path, mode="rb")
            total_size = path.stat().st_size
            # Also open a low-level fd for atomic positional reads
            # (os.pread is the only POSIX read primitive that doesn't move
            # the kernel cursor, which is what makes concurrent
            # read_chunk_at calls on the same handle safe). Tolerate
            # failure here — read_chunk still works via the aiofiles
            # cursor; only the random-access fast path needs the fd.
            random_access_fd = None
            try:
                random_access_fd = os.open(str(path), os.O_RDONLY)
            except OSError as fd_err:
                logger.warning(
                    f"Could not open random-access fd for {path}: {fd_err}; "
                    "read_chunk_at will be unavailable for this handle."
                )
            self._handles[handle_id] = handle
            self._file_info[handle_id] = {
                "path": path,
                "size": total_size,
                "bytes_read": 0,
                "created_at": time.time(),
                "random_access_fd": random_access_fd,
            }
            logger.debug(
                f"Opened file for read: {path} (handle_id={handle_id}, size={total_size} bytes)"
            )
            return {"success": True, "handle_id": handle_id, "total_size": total_size}
        except Exception as e:
            logger.error(f"Failed to open file for read: {path}, error: {e}")
            return {"success": False, "error": str(e)}

    @tool
    async def read_chunk_at(self, handle_id: str, offset: int, size: int = 2 * 1024 * 1024):
        """Atomic positional read — doesn't move the file cursor.

        Use this when the client wants to fire multiple concurrent reads
        on the same handle (e.g. to download a file in parallel slices).
        Built on os.pread, which is a single atomic POSIX syscall so
        overlapping calls don't step on each other.

        Args:
            handle_id: File handle ID from open_file_for_read.
            offset: Byte offset from the start of the file.
            size: Number of bytes to read. Default: 2 MiB.

        Returns:
            dict: {"success": True, "data": str (base64), "bytes_read": int,
                   "offset": int, "eof": bool} on success.
                  {"success": False, "error": str} on failure.
                  ``eof`` is True when the chunk reaches or passes the file
                  end (or returned an empty read).
        """
        if handle_id not in self._file_info:
            return {"success": False, "error": "Handle not found"}

        info = self._file_info[handle_id]
        fd = info.get("random_access_fd")
        if fd is None:
            return {
                "success": False,
                "error": "Handle has no random-access fd (file opened with older worker)",
            }

        if offset < 0:
            return {"success": False, "error": "offset must be >= 0"}
        if size <= 0:
            return {"success": False, "error": "size must be > 0"}

        try:
            data = await asyncio.to_thread(os.pread, fd, size, offset)
            bytes_read = len(data)
            total_size = info["size"]
            return {
                "success": True,
                "data": base64.b64encode(data).decode("utf-8"),
                "bytes_read": bytes_read,
                "offset": offset,
                "eof": bytes_read == 0 or (offset + bytes_read) >= total_size,
            }
        except Exception as e:
            logger.error(
                f"Failed read_chunk_at (handle_id={handle_id}, offset={offset}, size={size}): {e}"
            )
            return {"success": False, "error": str(e)}

    @tool
    async def stream_read(
        self,
        file_path: str,
        reply_to: str,
        ack_subject: str,
        offset: int = 0,
        chunk_size: int = 64 * 1024,
        window: int = 192,
    ):
        """Stream a file by PUSHING binary chunks to the caller's inbox.

        Unlike read_chunk_at (one request → one reply), this opens the file and
        spawns a background task that publishes a sequence of small binary
        chunks straight to ``reply_to``, flushing each. That removes the
        per-chunk caller↔agent round-trip (the throughput limiter in the
        request/reply path) while keeping every wire message small and
        one-in-flight — the only pattern that doesn't jam the
        Modal→DO-LB→nginx→NATS egress (≥128 KiB or concurrency → ~6 MiB freeze;
        see fileManager.ts streamFileToDisk for the measurements).

        Wire format per chunk message (RAW binary — no base64/JSON, so no 33%
        bloat):
            bytes[0:4] = big-endian uint32 sequence number (0-based)
            bytes[4]   = flags: bit0 = EOF (no payload), bit1 = ERROR
            bytes[5:]  = raw file bytes (absent for EOF / ERROR markers)

        Flow control: the caller publishes an ack (big-endian uint32 = highest
        contiguous seq received) to ``ack_subject`` every ~1 MiB. The streamer
        blocks if it gets ``window`` chunks ahead of the last ack, bounding
        unacked data to window*chunk_size (<< NATS max_pending 256 MiB). Ack
        value 0xFFFFFFFF cancels. No ack for 15 s while blocked → abort
        (self-cleaning; no orphan tasks).

        Returns immediately with {"success": True, "total_size": int}; bytes
        flow over ``reply_to`` independently.
        """
        if ".." in file_path:
            return {"success": False, "error": "File path cannot contain '..'"}
        path = self._get_root() / file_path
        if not path.exists():
            return {"success": False, "error": "File does not exist"}
        if not path.is_file():
            return {"success": False, "error": "Not a regular file"}

        # Clamp so every wire message stays small (≤ ~64 KiB) and the window
        # can't pin an unreasonable amount of unacked data.
        chunk_size = max(4096, min(int(chunk_size), 64 * 1024))
        window = max(8, min(int(window), 1024))

        try:
            total_size = path.stat().st_size
            fd = os.open(str(path), os.O_RDONLY)
        except Exception as e:
            return {"success": False, "error": str(e)}

        nc = self.worker.nc if getattr(self, "worker", None) else None
        if nc is None:
            os.close(fd)
            return {"success": False, "error": "No NATS connection available for streaming"}

        # Register a handle in _file_info so the chatroom's idle-reaper sees an
        # active transfer (has_active_tasks counts recent _file_info entries) and
        # won't tear the sandbox down mid-download. _stream_push removes it.
        handle_id = str(uuid.uuid4())
        self._file_info[handle_id] = {
            "path": path,
            "size": total_size,
            "bytes_read": 0,
            "created_at": time.time(),
            "random_access_fd": fd,
            "streaming": True,
        }

        asyncio.create_task(
            self._stream_push(
                nc, fd, handle_id, total_size, reply_to, ack_subject,
                max(0, int(offset)), chunk_size, window, str(path),
            )
        )
        return {"success": True, "total_size": total_size}

    async def _stream_push(self, nc, fd, handle_id, total_size, reply_to, ack_subject,
                           offset, chunk_size, window, path_str):
        """Background pusher for stream_read. Owns ``fd`` + the ack sub + the
        ``_file_info`` handle; always cleans up all three. Publishes raw binary
        chunks to ``reply_to`` with windowed flow control driven by acks on
        ``ack_subject``."""
        CANCEL = 0xFFFFFFFF
        ACK_TIMEOUT = 15.0
        state = {"acked": -1, "cancel": False}
        ev = asyncio.Event()

        async def on_ack(msg):
            try:
                if len(msg.data) >= 4:
                    s = struct.unpack(">I", msg.data[:4])[0]
                    if s == CANCEL:
                        state["cancel"] = True
                    elif s > state["acked"]:
                        state["acked"] = s
            except Exception:
                pass
            ev.set()

        sub = await nc.subscribe(ack_subject, cb=on_ack)
        seq = 0
        pos = offset
        sent_bytes = 0
        try:
            while pos < total_size:
                if state["cancel"]:
                    logger.info(f"[stream_read] client cancelled: {path_str}")
                    return
                # Windowed flow control: never get > window chunks ahead of ack.
                while (seq - state["acked"] - 1) >= window and not state["cancel"]:
                    ev.clear()
                    try:
                        await asyncio.wait_for(ev.wait(), ACK_TIMEOUT)
                    except asyncio.TimeoutError:
                        logger.warning(f"[stream_read] ack timeout, aborting: {path_str}")
                        return
                if state["cancel"]:
                    return
                n = min(chunk_size, total_size - pos)
                data = await asyncio.to_thread(os.pread, fd, n, pos)
                if not data:
                    break
                # RAW binary message: 5-byte header + payload. One flush per
                # chunk keeps ≤1 small message in flight (jam-free egress).
                await nc.publish(reply_to, struct.pack(">IB", seq, 0) + data)
                await nc.flush()
                pos += len(data)
                sent_bytes += len(data)
                seq += 1
                # Keep the idle-reaper's transfer_handles window fresh on long
                # downloads (it only counts handles < 3600s old).
                if seq % 256 == 0:
                    info = self._file_info.get(handle_id)
                    if info is not None:
                        info["created_at"] = time.time()
                        info["bytes_read"] = sent_bytes
            if not state["cancel"]:
                await nc.publish(reply_to, struct.pack(">IB", seq, 1))  # EOF
                await nc.flush()
                logger.debug(f"[stream_read] done {path_str}: {sent_bytes} bytes in {seq} chunks")
        except Exception as e:
            logger.error(f"[stream_read] error streaming {path_str}: {e}")
            try:
                await nc.publish(reply_to, struct.pack(">IB", seq, 2))  # ERROR
                await nc.flush()
            except Exception:
                pass
        finally:
            try:
                await sub.unsubscribe()
            except Exception:
                pass
            self._file_info.pop(handle_id, None)
            try:
                os.close(fd)
            except Exception:
                pass

    @tool
    async def read_chunk(self, handle_id: str, size: int = 512 * 1024):
        """Read a chunk from an open file handle (async).

        Args:
            handle_id: File handle ID from open_file_for_read.
            size: Number of bytes to read (before base64 encoding). Default: 512KB.

        Returns:
            dict: {"success": True, "data": str (base64), "bytes_read": int, "eof": bool}
                  {"success": False, "error": str} on failure
        """
        if handle_id not in self._handles:
            return {"success": False, "error": "Handle not found"}

        handle = self._handles[handle_id]
        info = self._file_info[handle_id]

        try:
            # Read data asynchronously
            data = await handle.read(size)

            if not data:
                return {"success": True, "data": "", "bytes_read": 0, "eof": True}

            bytes_read = len(data)
            info["bytes_read"] += bytes_read

            return {
                "success": True,
                "data": base64.b64encode(data).decode("utf-8"),
                "bytes_read": bytes_read,
                "eof": bytes_read < size,
            }
        except Exception as e:
            logger.error(f"Failed to read chunk (handle_id={handle_id}): {e}")
            return {"success": False, "error": str(e)}

    @tool
    async def read_file(
        self, file_path: str, receive_chunk=None, chunk_size: int = 1024
    ):
        """Read a file (streaming or non-streaming mode).

        Args:
            file_path: Relative path to the file (cannot contain '..')
            receive_chunk: Optional callback for streaming mode (direct connections)
            chunk_size: Chunk size for streaming mode. Default: 1KB.

        Returns:
            Non-streaming: {"success": True, "data": str (base64), "total_size": int, "encoding": "base64"}
            Streaming: {"success": True}
            Error: {"success": False, "error": str}
        """
        if ".." in file_path:
            return {"success": False, "error": "File path cannot contain '..'"}

        path = self._get_root() / file_path
        if not path.exists():
            return {"success": False, "error": "File does not exist"}

        if receive_chunk is None:
            # Non-streaming mode: return full file content (base64 encoded for JSON compatibility)
            try:
                async with aiofiles.open(path, mode="rb") as f:
                    file_data = await f.read()
                    return {
                        "success": True,
                        "data": base64.b64encode(file_data).decode("utf-8"),
                        "total_size": len(file_data),
                        "encoding": "base64",
                    }
            except Exception as e:
                logger.error(f"Failed to read file: {path}, error: {e}")
                return {"success": False, "error": str(e)}
        else:
            # Streaming mode: use callback function for direct connections
            try:
                async with aiofiles.open(path, mode="rb") as f:
                    while True:
                        data = await f.read(chunk_size)
                        if not data:
                            break
                        await receive_chunk(data)
                return {"success": True}
            except Exception as e:
                logger.error(f"Failed to stream file: {path}, error: {e}")
                return {"success": False, "error": str(e)}

