from __future__ import annotations

"""Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Filesystem-backed transfer session store with resumable upload/download staging.
"""

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from json import dump, load
from os import listdir, makedirs, remove
from os.path import abspath, commonpath, exists, getsize, join, realpath
from uuid import uuid4

from mcp_proxy_adapter.transfer.checksums import (
    compute_file_checksum,
    verify_file_checksum,
)
from mcp_proxy_adapter.transfer.compression import (
    compress_file,
    decompress_file,
    validate_compression_algorithm,
)
from mcp_proxy_adapter.transfer.errors import (
    TransferAckError,
    TransferChecksumMismatchError,
    TransferCompressionError,
    TransferError,
    TransferExpiredError,
    TransferRangeNotSatisfiableError,
    TransferResumeConflictError,
    TransferSessionNotFoundError,
    TransferTooLargeError,
)
from mcp_proxy_adapter.transfer.models import (
    TransferDownloadReceipt,
    TransferReceipt,
    TransferSession,
    TransferStatus,
)


@dataclass(frozen=True)
class ExpireStaleSessionsResult:
    """Outcome of ``TransferServerStore.expire_stale_sessions``."""

    count: int
    expired: tuple[tuple[str, str], ...]
    """``(transfer_id, direction)`` for each session newly marked expired."""


class TransferServerStore:
    def __init__(
        self,
        storage_dir: str,
        session_ttl_seconds: int,
        max_file_size_bytes: int,
        default_chunk_size_bytes: int,
        max_chunk_size_bytes: int,
        server_gzip_download_buffer_dir: str | None = None,
    ) -> None:
        if session_ttl_seconds < 60:
            raise ValueError("session_ttl_seconds must be >= 60")
        if max_file_size_bytes < 1:
            raise ValueError("max_file_size_bytes must be >= 1")
        if default_chunk_size_bytes < 1:
            raise ValueError("default_chunk_size_bytes must be >= 1")
        if max_chunk_size_bytes < default_chunk_size_bytes:
            raise ValueError("max_chunk_size_bytes must be >= default_chunk_size_bytes")
        self._storage_dir = abspath(storage_dir)
        self._gzip_buffer_dir = abspath(
            server_gzip_download_buffer_dir
            if server_gzip_download_buffer_dir
            else self._storage_dir
        )
        self._session_ttl_seconds = session_ttl_seconds
        self._max_file_size_bytes = max_file_size_bytes
        self._default_chunk_size_bytes = default_chunk_size_bytes
        self._max_chunk_size_bytes = max_chunk_size_bytes
        makedirs(self._sessions_dir(), exist_ok=True)
        makedirs(self._files_dir(), exist_ok=True)
        makedirs(self._gzip_buffer_dir, exist_ok=True)

    def _gzip_buffer_root(self) -> str:
        return abspath(realpath(self._gzip_buffer_dir))

    def _path_is_under_gzip_buffer(self, path: str) -> bool:
        try:
            p = abspath(realpath(path))
            root = self._gzip_buffer_root()
        except OSError:
            return False
        try:
            return commonpath([root, p]) == root
        except ValueError:
            return False

    def _unlink_gzip_artifact_if_owned(self, session: TransferSession) -> None:
        artifact = session.compressed_storage_path
        if not artifact:
            return
        if not self._path_is_under_gzip_buffer(artifact):
            return
        if exists(artifact):
            remove(artifact)

    def _sessions_dir(self) -> str:
        return join(self._storage_dir, "sessions")

    def _files_dir(self) -> str:
        return join(self._storage_dir, "files")

    def _metadata_path(self, transfer_id: str) -> str:
        return join(self._sessions_dir(), f"{transfer_id}.json")

    def _new_transfer_id(self) -> str:
        return f"tr_{uuid4().hex}"

    def _parse_iso_utc(self, value: str) -> datetime:
        normalized = value.replace("Z", "+00:00") if value.endswith("Z") else value
        return datetime.fromisoformat(normalized)

    def _utc_now_iso(self) -> str:
        return datetime.now(timezone.utc).isoformat()

    def _is_terminal_status(self, status: str) -> bool:
        return status in (
            TransferStatus.FAILED.value,
            TransferStatus.EXPIRED.value,
            TransferStatus.UPLOADED.value,
            TransferStatus.COMPLETED.value,
            TransferStatus.CONSUMED.value,
        )

    def _session_is_expired(self, session: TransferSession, *, now: datetime) -> bool:
        if session.expires_at is None:
            return False
        return self._parse_iso_utc(session.expires_at) < now

    def try_load_session_metadata(self, transfer_id: str) -> TransferSession | None:
        """
        Load session JSON without expiry / consumed guards.

        Used for WebSocket correlation envelopes and idempotent ack replay.
        """
        return self._load_session_raw(str(transfer_id).strip())

    def _load_session_raw(self, transfer_id: str) -> TransferSession | None:
        path = self._metadata_path(transfer_id)
        if not exists(path):
            return None
        try:
            with open(path, encoding="utf-8") as f:
                raw = load(f)
            return TransferSession.from_dict(raw)
        except ValueError as exc:
            raise TransferError(
                message="Corrupt or invalid session metadata",
                transfer_id=transfer_id,
                phase="load",
            ) from exc

    def _persist_session(self, session: TransferSession) -> None:
        with open(session.metadata_path, "w", encoding="utf-8") as f:
            dump(session.to_dict(), f, indent=2)

    def _validate_checksum_value(self, checksum_value: str) -> None:
        if not isinstance(checksum_value, str) or not checksum_value:
            raise TransferError(message="Invalid checksum_value", phase="create")
        if len(checksum_value) != 64:
            raise TransferError(message="Invalid checksum_value", phase="create")
        for c in checksum_value:
            if c not in "0123456789abcdef":
                raise TransferError(message="Invalid checksum_value", phase="create")

    def _upload_receipt_from_session(
        self, session: TransferSession, *, completed: bool
    ) -> TransferReceipt:
        return TransferReceipt(
            transfer_id=session.transfer_id,
            filename=session.filename,
            size_bytes=session.size_bytes,
            checksum_algorithm=session.checksum_algorithm,
            checksum_value=session.checksum_value,
            compression=session.compression,
            chunk_size=session.chunk_size,
            offset=session.offset,
            status=session.status,
            completed=completed,
            expires_at=session.expires_at,
            plaintext_size_bytes=session.plaintext_size_bytes,
        )

    @staticmethod
    def _optional_correlation_str(value: object | None) -> str | None:
        if value is None:
            return None
        s = str(value).strip()
        return s or None

    def _download_source_deletable(self, session: TransferSession) -> bool:
        if not session.delete_source_on_ack:
            return False
        path = abspath(realpath(session.storage_path))
        try:
            files_root = abspath(realpath(self._files_dir()))
            gzip_root = self._gzip_buffer_root()
            if commonpath([files_root, path]) == files_root:
                return True
            if commonpath([gzip_root, path]) == gzip_root:
                return True
        except (OSError, ValueError):
            return False
        return False

    def create_upload_session(
        self,
        filename: str,
        size_bytes: int,
        checksum_algorithm: str,
        checksum_value: str,
        compression: str,
        chunk_size: int | None = None,
        *,
        job_id: str | None = None,
        correlation_id: str | None = None,
        command: str | None = None,
        inbound_transfer_id: str | None = None,
    ) -> TransferReceipt:
        validate_compression_algorithm(compression)
        if checksum_algorithm != "sha256":
            raise TransferError(
                message='checksum_algorithm must be "sha256"',
                phase="create",
            )
        self._validate_checksum_value(checksum_value)
        if size_bytes > self._max_file_size_bytes:
            raise TransferTooLargeError(
                message="File size exceeds configured maximum",
                phase="create",
            )
        resolved_chunk = (
            chunk_size if chunk_size is not None else self._default_chunk_size_bytes
        )
        if resolved_chunk <= 0 or resolved_chunk > self._max_chunk_size_bytes:
            raise TransferError(
                message="chunk_size out of allowed bounds", phase="create"
            )
        transfer_id = self._new_transfer_id()
        metadata_path = self._metadata_path(transfer_id)
        storage_path = join(self._files_dir(), f"{transfer_id}.upload")
        open(storage_path, "wb").close()
        expires_at = (
            datetime.now(timezone.utc) + timedelta(seconds=self._session_ttl_seconds)
        ).isoformat()
        session = TransferSession(
            transfer_id=transfer_id,
            filename=filename,
            direction="upload",
            size_bytes=size_bytes,
            checksum_value=checksum_value,
            chunk_size=resolved_chunk,
            status=TransferStatus.CREATED.value,
            storage_path=storage_path,
            metadata_path=metadata_path,
            checksum_algorithm=checksum_algorithm,
            compression=compression,
            offset=0,
            expires_at=expires_at,
            error_message=None,
            compressed_storage_path=None,
            plaintext_size_bytes=None,
            job_id=self._optional_correlation_str(job_id),
            correlation_id=self._optional_correlation_str(correlation_id),
            command=self._optional_correlation_str(command),
            inbound_transfer_id=self._optional_correlation_str(inbound_transfer_id),
            delete_source_on_ack=False,
        )
        self._persist_session(session)
        return self._upload_receipt_from_session(session, completed=False)

    def write_upload_chunk(
        self, transfer_id: str, offset: int, chunk_bytes: bytes
    ) -> TransferReceipt:
        session = self.get_session(transfer_id)
        if session.direction != "upload":
            raise TransferError(
                message="Not an upload session",
                transfer_id=transfer_id,
                phase="upload",
            )
        if self._is_terminal_status(session.status):
            raise TransferError(
                message="Upload session is not active",
                transfer_id=transfer_id,
                phase="upload",
            )
        if offset != session.offset:
            raise TransferResumeConflictError(
                message="Upload offset mismatch",
                transfer_id=transfer_id,
                expected_offset=session.offset,
                actual_offset=offset,
            )
        if len(chunk_bytes) > self._max_chunk_size_bytes:
            raise TransferTooLargeError(
                message="Chunk exceeds max_chunk_size_bytes",
                transfer_id=transfer_id,
                phase="upload",
            )
        if session.offset + len(chunk_bytes) > session.size_bytes:
            raise TransferTooLargeError(
                message="Upload would exceed declared size_bytes",
                transfer_id=transfer_id,
                phase="upload",
            )
        with open(session.storage_path, "ab") as staged:
            staged.write(chunk_bytes)
        session.offset += len(chunk_bytes)
        if session.offset < session.size_bytes:
            session.status = TransferStatus.UPLOADING.value
        elif session.offset == session.size_bytes:
            session.status = TransferStatus.UPLOADING.value
        self._persist_session(session)
        return self._upload_receipt_from_session(session, completed=False)

    def complete_upload_session(self, transfer_id: str) -> TransferReceipt:
        session = self.get_session(transfer_id)
        if session.direction != "upload":
            raise TransferError(
                message="Not an upload session",
                transfer_id=transfer_id,
                phase="complete",
            )
        if session.status == TransferStatus.UPLOADED.value:
            return self._upload_receipt_from_session(session, completed=True)
        if session.status in (
            TransferStatus.FAILED.value,
            TransferStatus.EXPIRED.value,
        ):
            raise TransferError(
                message="Session is not active",
                transfer_id=transfer_id,
                phase="complete",
            )
        if getsize(session.storage_path) != session.size_bytes:
            session.status = TransferStatus.FAILED.value
            session.error_message = "Staged size mismatch"
            self._persist_session(session)
            raise TransferError(
                message="Staged size mismatch",
                transfer_id=transfer_id,
                phase="complete",
            )
        if session.offset != session.size_bytes:
            session.status = TransferStatus.FAILED.value
            session.error_message = "Incomplete upload"
            self._persist_session(session)
            raise TransferError(
                message="Incomplete upload",
                transfer_id=transfer_id,
                phase="complete",
            )
        try:
            if session.compression == "identity":
                verify_file_checksum(
                    session.storage_path,
                    session.checksum_value,
                    transfer_id=transfer_id,
                    phase="verify",
                )
            elif session.compression == "gzip":
                verify_tmp = join(self._files_dir(), f"{transfer_id}.verify_tmp")
                try:
                    decompress_file(session.storage_path, verify_tmp, "gzip")
                    verify_file_checksum(
                        verify_tmp,
                        session.checksum_value,
                        transfer_id=transfer_id,
                        phase="verify",
                    )
                finally:
                    if exists(verify_tmp):
                        remove(verify_tmp)
        except (TransferChecksumMismatchError, TransferCompressionError) as exc:
            session.status = TransferStatus.FAILED.value
            session.error_message = str(exc)
            self._persist_session(session)
            raise
        session.status = TransferStatus.UPLOADED.value
        session.offset = session.size_bytes
        session.error_message = None
        self._persist_session(session)
        return self._upload_receipt_from_session(session, completed=True)

    def create_download_session(
        self,
        source_path: str,
        filename: str,
        checksum_algorithm: str,
        compression: str,
        chunk_size: int | None = None,
        *,
        job_id: str | None = None,
        correlation_id: str | None = None,
        command: str | None = None,
        inbound_transfer_id: str | None = None,
        delete_source_on_ack: bool = False,
    ) -> TransferReceipt:
        abs_src = abspath(source_path)
        if not exists(abs_src):
            raise TransferError(message="Source file does not exist", phase="download")
        validate_compression_algorithm(compression)
        if checksum_algorithm != "sha256":
            raise TransferError(
                message='checksum_algorithm must be "sha256"',
                phase="download",
            )
        plaintext_size = getsize(abs_src)
        if plaintext_size > self._max_file_size_bytes:
            raise TransferTooLargeError(
                message="File size exceeds configured maximum",
                phase="download",
            )
        checksum_value = compute_file_checksum(abs_src)
        resolved_chunk = (
            chunk_size if chunk_size is not None else self._default_chunk_size_bytes
        )
        if resolved_chunk <= 0 or resolved_chunk > self._max_chunk_size_bytes:
            raise TransferError(
                message="chunk_size out of allowed bounds", phase="download"
            )
        transfer_id = self._new_transfer_id()
        metadata_path = self._metadata_path(transfer_id)
        expires_at = (
            datetime.now(timezone.utc) + timedelta(seconds=self._session_ttl_seconds)
        ).isoformat()
        compressed_storage_path: str | None = None
        wire_size: int
        plaintext_size_bytes: int | None
        if compression == "gzip":
            artifact_path = join(self._gzip_buffer_dir, f"{transfer_id}.download.gzip")
            try:
                compress_file(abs_src, artifact_path, "gzip")
            except TransferCompressionError:
                if exists(artifact_path):
                    remove(artifact_path)
                raise
            compressed_storage_path = artifact_path
            wire_size = getsize(artifact_path)
            plaintext_size_bytes = plaintext_size
        else:
            wire_size = plaintext_size
            plaintext_size_bytes = None
        session = TransferSession(
            transfer_id=transfer_id,
            filename=filename,
            direction="download",
            size_bytes=wire_size,
            checksum_value=checksum_value,
            chunk_size=resolved_chunk,
            status=TransferStatus.CREATED.value,
            storage_path=abs_src,
            metadata_path=metadata_path,
            checksum_algorithm=checksum_algorithm,
            compression=compression,
            offset=0,
            expires_at=expires_at,
            error_message=None,
            compressed_storage_path=compressed_storage_path,
            plaintext_size_bytes=plaintext_size_bytes,
            job_id=self._optional_correlation_str(job_id),
            correlation_id=self._optional_correlation_str(correlation_id),
            command=self._optional_correlation_str(command),
            inbound_transfer_id=self._optional_correlation_str(inbound_transfer_id),
            delete_source_on_ack=bool(delete_source_on_ack),
        )
        self._persist_session(session)
        return self._upload_receipt_from_session(session, completed=False)

    def acknowledge_download_transfer(
        self,
        transfer_id: str,
        *,
        checksum_algorithm: str,
        checksum_value: str,
        job_id: str | None = None,
    ) -> dict[str, object]:
        """
        Idempotent post-download ack: verify checksum, optional job binding,
        delete staging when policy allows, mark session consumed.
        """
        tid = str(transfer_id).strip()
        if not tid:
            raise TransferAckError(
                message="transfer_id is required",
                transfer_id=None,
                phase="ack",
            )
        session = self._load_session_raw(tid)
        if session is None:
            raise TransferSessionNotFoundError(
                message="Transfer session not found",
                transfer_id=tid,
                phase="ack",
            )
        if session.direction != "download":
            raise TransferAckError(
                message="transfer_ack applies only to download sessions",
                transfer_id=tid,
                phase="ack",
            )
        if session.status == TransferStatus.CONSUMED.value:
            return {
                "success": True,
                "transfer_id": tid,
                "idempotent": True,
            }
        if session.status != TransferStatus.COMPLETED.value:
            raise TransferAckError(
                message="Download is not complete; finish HTTP transfer before ack",
                transfer_id=tid,
                phase="ack",
            )
        algo = str(checksum_algorithm).strip().lower()
        if algo != "sha256":
            raise TransferAckError(
                message='checksum_algorithm must be "sha256"',
                transfer_id=tid,
                phase="ack",
            )
        self._validate_checksum_value(checksum_value)
        digest = str(checksum_value).strip().lower()
        expected = str(session.checksum_value).strip().lower()
        if digest != expected:
            raise TransferChecksumMismatchError(
                message="Client checksum does not match session metadata",
                transfer_id=tid,
                checksum_expected=expected,
                checksum_actual=digest,
                phase="ack",
            )
        bound = self._optional_correlation_str(job_id)
        if session.job_id and bound and session.job_id != bound:
            raise TransferAckError(
                message="job_id does not match session binding",
                transfer_id=tid,
                phase="ack",
            )
        if session.job_id and bound is None:
            raise TransferAckError(
                message="job_id is required for this transfer session",
                transfer_id=tid,
                phase="ack",
            )
        self._unlink_gzip_artifact_if_owned(session)
        if self._download_source_deletable(session):
            sp = abspath(session.storage_path)
            if exists(sp):
                try:
                    remove(sp)
                except OSError as exc:
                    raise TransferAckError(
                        message=f"Failed to remove staged file: {exc}",
                        transfer_id=tid,
                        phase="ack",
                    ) from exc
        session.compressed_storage_path = None
        session.status = TransferStatus.CONSUMED.value
        session.error_message = None
        self._persist_session(session)
        return {
            "success": True,
            "transfer_id": tid,
            "idempotent": False,
        }

    def read_download_chunk(
        self, transfer_id: str, offset: int, max_bytes: int
    ) -> tuple[bytes, TransferDownloadReceipt]:
        session = self.get_session(transfer_id)
        if session.direction != "download":
            raise TransferError(
                message="Not a download session",
                transfer_id=transfer_id,
                phase="download",
            )
        if session.compression == "gzip":
            read_path = session.compressed_storage_path
            if not read_path:
                raise TransferError(
                    message="Gzip download session missing compressed artifact path",
                    transfer_id=transfer_id,
                    phase="download",
                )
        else:
            read_path = session.storage_path
        if offset > session.size_bytes:
            raise TransferRangeNotSatisfiableError(
                message="Offset beyond end of file",
                transfer_id=transfer_id,
                expected_offset=session.size_bytes,
                actual_offset=offset,
            )
        if offset != session.offset:
            raise TransferResumeConflictError(
                message="Download offset mismatch",
                transfer_id=transfer_id,
                expected_offset=session.offset,
                actual_offset=offset,
            )
        read_len = min(max_bytes, session.size_bytes - offset)
        with open(read_path, "rb") as src:
            src.seek(offset)
            data = src.read(read_len)
        new_offset = offset + len(data)
        session.offset = new_offset
        if new_offset == session.size_bytes:
            session.status = TransferStatus.COMPLETED.value
            if session.compression == "gzip":
                self._unlink_gzip_artifact_if_owned(session)
                session.compressed_storage_path = None
        else:
            session.status = TransferStatus.DOWNLOADING.value
        self._persist_session(session)
        receipt = TransferDownloadReceipt(
            transfer_id=session.transfer_id,
            destination_path="",
            offset=new_offset,
            size_bytes=session.size_bytes,
            checksum_algorithm=session.checksum_algorithm,
            checksum_value=session.checksum_value,
            compression=session.compression,
            status=session.status,
            completed=(new_offset == session.size_bytes),
            checksum_verified=False,
        )
        return (data, receipt)

    def get_session(self, transfer_id: str) -> TransferSession:
        session = self._load_session_raw(transfer_id)
        if session is None:
            raise TransferSessionNotFoundError(
                message="Transfer session not found",
                transfer_id=transfer_id,
            )
        now = datetime.now(timezone.utc)
        if self._session_is_expired(session, now=now) and not self._is_terminal_status(
            session.status
        ):
            raise TransferExpiredError(
                message="Transfer session expired",
                transfer_id=transfer_id,
            )
        if session.status == TransferStatus.EXPIRED.value:
            raise TransferExpiredError(
                message="Transfer session expired",
                transfer_id=transfer_id,
            )
        if session.status == TransferStatus.CONSUMED.value:
            raise TransferSessionNotFoundError(
                message="Transfer session not found",
                transfer_id=transfer_id,
                phase="lookup",
            )
        return session

    def get_committed_upload_path(self, transfer_id: str) -> str:
        session = self.get_session(transfer_id)
        if session.direction != "upload":
            raise TransferError(
                message="Not an upload session",
                transfer_id=transfer_id,
                phase="commit",
            )
        if session.status != TransferStatus.UPLOADED.value:
            raise TransferError(
                message="Upload not complete",
                transfer_id=transfer_id,
                phase="commit",
            )
        return session.storage_path

    def get_completed_transfer(self, transfer_id: str) -> dict[str, object]:
        """Return completed upload metadata for downstream server-side processing."""
        session = self.get_session(transfer_id)
        if session.direction != "upload":
            raise TransferError(
                message="Not an upload session",
                transfer_id=transfer_id,
                phase="commit",
            )
        if session.status != TransferStatus.UPLOADED.value:
            raise TransferError(
                message="Upload not complete",
                transfer_id=transfer_id,
                phase="commit",
            )
        return {
            "transfer_id": session.transfer_id,
            "local_path": session.storage_path,
            "filename": session.filename,
            "compression": session.compression,
            "checksum_algorithm": session.checksum_algorithm,
            "checksum_value": session.checksum_value,
            "size_bytes": session.size_bytes,
            "chunk_size": session.chunk_size,
            "offset": session.offset,
            "status": session.status,
            "expires_at": session.expires_at,
        }

    def expire_stale_sessions(
        self, now_iso: str | None = None
    ) -> ExpireStaleSessionsResult:
        now = (
            datetime.fromisoformat(now_iso.replace("Z", "+00:00"))
            if now_iso
            else datetime.now(timezone.utc)
        )
        expired_rows: list[tuple[str, str]] = []
        sessions_dir = self._sessions_dir()
        for name in listdir(sessions_dir):
            if not name.endswith(".json"):
                continue
            path = join(sessions_dir, name)
            try:
                with open(path, encoding="utf-8") as f:
                    raw = load(f)
                session = TransferSession.from_dict(raw)
            except (OSError, ValueError, TypeError) as exc:
                raise TransferError(
                    message="Corrupt session metadata", phase="expire"
                ) from exc
            if session.status == TransferStatus.CONSUMED.value:
                if session.expires_at is None:
                    continue
                if not self._session_is_expired(session, now=now):
                    continue
                try:
                    remove(path)
                except OSError:
                    pass
                continue
            if self._is_terminal_status(session.status):
                continue
            if session.expires_at is None:
                continue
            if not self._session_is_expired(session, now=now):
                continue
            tid = str(session.transfer_id).strip()
            direction = str(session.direction).strip()
            self._unlink_gzip_artifact_if_owned(session)
            if (
                session.direction == "upload"
                and exists(session.storage_path)
                and session.storage_path.startswith(self._files_dir())
            ):
                remove(session.storage_path)
            session.compressed_storage_path = None
            session.status = TransferStatus.EXPIRED.value
            session.error_message = None
            self._persist_session(session)
            if tid:
                expired_rows.append((tid, direction))
        return ExpireStaleSessionsResult(
            count=len(expired_rows), expired=tuple(expired_rows)
        )
