#!/usr/bin/env python3
"""
mcp_drive_gccl_compress.py — Download a Google Drive folder via MCP, GCCL-compress it, upload archive.

Workflow:
  1. Recursively list all files in a Drive folder using the Google Drive MCP server.
  2. Download each file to a local temp clone.
  3. For each file, attempt GCCL delta compression using a previously-seen similar file as reference.
  4. Pack compressed results into a .tar.zst archive with a JSON manifest.
  5. Upload the archive back to Google Drive via MCP.
  6. Optionally move the original folder to trash via MCP.

Usage:
  python3 mcp_drive_gccl_compress.py <drive-folder-id> [options]

No rclone is used — only the Google Drive MCP server and gccl_waveprobe.
"""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Insert Research Stack shim path for gccl_waveprobe
REPO = Path("/home/allaun/Research Stack")
SHIM_DIR = REPO / "4-Infrastructure/shim"
sys.path.insert(0, str(SHIM_DIR))

try:
    import gccl_waveprobe as gw
except ImportError as e:
    raise RuntimeError(f"Cannot import gccl_waveprobe from {SHIM_DIR}: {e}")


DEFAULT_MCP_COMMAND = ["npx", "-y", "@piotr-agier/google-drive-mcp"]


class DriveMcpClient:
    """Very thin JSON-RPC stdio client for the Google Drive MCP server."""

    def __init__(self, env: Optional[Dict[str, str]] = None):
        self._proc = subprocess.Popen(
            DEFAULT_MCP_COMMAND,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=False,
            env={**os.environ, **(env or {})},
        )
        self._next_id = 1
        self._initialize()

    def _initialize(self) -> None:
        self._call(
            "initialize",
            {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "mcp-drive-gccl", "version": "1.0.0"},
            },
        )

    def _send(self, msg: dict) -> None:
        payload = json.dumps(msg).encode("utf-8") + b"\n"
        self._proc.stdin.write(payload)  # type: ignore[union-attr]
        self._proc.stdin.flush()  # type: ignore[union-attr]

    def _recv(self) -> dict:
        line = self._proc.stdout.readline()  # type: ignore[union-attr]
        if not line:
            raise RuntimeError("MCP server closed stdout")
        return json.loads(line.decode("utf-8"))

    def _call(self, method: str, params: dict) -> dict:
        msg = {"jsonrpc": "2.0", "id": self._next_id, "method": method, "params": params}
        self._next_id += 1
        self._send(msg)
        while True:
            resp = self._recv()
            if resp.get("id") == msg["id"]:
                if "error" in resp:
                    raise RuntimeError(f"MCP error: {resp['error']}")
                return resp["result"]

    def list_folder(
        self,
        folder_id: Optional[str] = None,
        page_size: int = 100,
        page_token: Optional[str] = None,
    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
        params: Dict[str, Any] = {"pageSize": page_size}
        if folder_id:
            params["folderId"] = folder_id
        if page_token:
            params["pageToken"] = page_token
        result = self._call("tools/call", {"name": "listFolder", "arguments": params})
        text = self._extract_text(result)
        items = self._parse_folder_listing(text)
        next_token = self._parse_next_page_token(text)
        return items, next_token

    def list_folder_all(self, folder_id: Optional[str] = None) -> List[Dict[str, Any]]:
        all_items: List[Dict[str, Any]] = []
        token: Optional[str] = None
        while True:
            items, token = self.list_folder(folder_id=folder_id, page_token=token)
            all_items.extend(items)
            if not token:
                break
        return all_items

    def download_file(self, file_id: str, local_path: Path) -> None:
        result = self._call(
            "tools/call",
            {"name": "downloadFile", "arguments": {"fileId": file_id, "localPath": str(local_path), "overwrite": True}},
        )
        text = self._extract_text(result)
        if "Downloaded" not in text and "Saved" not in text:
            raise RuntimeError(f"Download failed: {text}")

    def upload_file(self, local_path: Path, parent_folder_id: Optional[str] = None) -> str:
        args: Dict[str, Any] = {"localPath": str(local_path)}
        if parent_folder_id:
            args["parentFolderId"] = parent_folder_id
        result = self._call("tools/call", {"name": "uploadFile", "arguments": args})
        text = self._extract_text(result)
        if "Uploaded" not in text and "uploaded" not in text:
            raise RuntimeError(f"Upload failed: {text}")
        return text

    def delete_item(self, item_id: str) -> None:
        self._call("tools/call", {"name": "deleteItem", "arguments": {"itemId": item_id}})

    def close(self) -> None:
        if self._proc.poll() is None:
            self._proc.stdin.close()  # type: ignore[union-attr]
            self._proc.wait(timeout=5)

    @staticmethod
    def _parse_next_page_token(text: str) -> Optional[str]:
        for line in text.splitlines():
            line = line.strip()
            if line.lower().startswith("nextpagetoken:") or line.lower().startswith("next page token:"):
                return line.split(":", 1)[1].strip()
            if "nextpagetoken" in line.lower() and ":" in line:
                return line.split(":", 1)[1].strip()
        return None

    @staticmethod
    def _extract_text(result: dict) -> str:
        for item in result.get("content", []):
            if item.get("type") == "text":
                return item.get("text", "")
        return ""

    @staticmethod
    def _parse_folder_listing(text: str) -> List[Dict[str, Any]]:
        items: List[Dict[str, Any]] = []
        for line in text.splitlines():
            line = line.strip()
            if not line or line.startswith("Contents"):
                continue
            # Format: 📄 name (mime) [id: ID, path: ...] [created: ..., modified: ...]
            # or     📁 name (ID: ID)
            if line.startswith("📄"):
                rest = line[2:].strip()
                # Format: "name (ID: FILE_ID, path: ...)" or "name (mime) (ID: ...)"
                if "(ID: " in rest:
                    name = rest.split("(ID: ")[0].strip()
                    # Remove trailing mime like "(text/plain)" from name
                    if name.endswith(")") and "(" in name:
                        name = name.rsplit("(", 1)[0].strip()
                    fid = rest.split("(ID: ")[1].split(",", 1)[0].rstrip(")").strip()
                    items.append({"type": "file", "name": name, "id": fid})
            elif line.startswith("📁"):
                # Format: name (ID: ID)
                rest = line[2:].strip()
                if "(ID: " in rest:
                    name = rest.split("(ID: ")[0].strip()
                    fid = rest.split("(ID: ")[1].rstrip(")").strip()
                    items.append({"type": "folder", "name": name, "id": fid})
        return items


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        while chunk := f.read(64 * 1024):
            h.update(chunk)
    return h.hexdigest()


def gccl_compress_file(
    src_path: Path,
    reference_path: Optional[Path] = None,
    threshold_q16: int = 32768,
    max_gccl_size: int = 100 * 1024 * 1024,
) -> Tuple[bytes, Dict[str, Any]]:
    """Compress a single file with GCCL. Optionally use reference for delta.
    Files larger than max_gccl_size are stored raw to avoid memory issues."""
    original_size = src_path.stat().st_size

    if original_size > max_gccl_size:
        return b"", {
            "mode": "raw-oversize",
            "reference": str(reference_path) if reference_path else None,
            "original_size": original_size,
            "compressed_size": original_size,
            "ratio": 1.0,
            "decision": "OVERSIZE",
        }

    if original_size == 0:
        return b"", {
            "mode": "raw",
            "reference": str(reference_path) if reference_path else None,
            "original_size": 0,
            "compressed_size": 0,
            "ratio": 1.0,
            "decision": "EMPTY",
        }

    # If the reference file is larger than max_gccl_size, do not use it to avoid MemoryError
    if reference_path and reference_path.stat().st_size > max_gccl_size:
        reference_path = None

    data = src_path.read_bytes()
    reference = reference_path.read_bytes() if reference_path else None

    result = gw.gccl_delta_compress(
        data,
        reference=reference,
        gccl_threshold_q16=threshold_q16,
        wave_probe=gw.WaveProbe(id=1, sample_rate=48000, buffer_size=256),
    )

    if result.gccl_decision == gw.Decision.ACCEPT and result.compressed_size < result.original_size:
        return result.compressed, {
            "mode": "delta" if reference else "raw-gccl",
            "reference": str(reference_path) if reference_path else None,
            "original_size": result.original_size,
            "compressed_size": result.compressed_size,
            "ratio": result.compression_ratio,
            "decision": str(result.gccl_decision),
        }
    else:
        # Compression rejected or expanded — store raw
        return data, {
            "mode": "raw",
            "reference": str(reference_path) if reference_path else None,
            "original_size": len(data),
            "compressed_size": len(data),
            "ratio": 1.0,
            "decision": str(result.gccl_decision) if result else "N/A",
        }


def download_folder_recursive(
    client: DriveMcpClient,
    folder_id: str,
    local_root: Path,
    rename_map: Optional[Dict[str, str]] = None,
) -> List[Dict[str, Any]]:
    """Recursively download a Drive folder. Returns manifest of downloaded files."""
    files: List[Dict[str, Any]] = []
    _download_folder_recursive(client, folder_id, local_root, files, rename_map or {})
    return files


def _download_folder_recursive(
    client: DriveMcpClient,
    folder_id: str,
    local_dir: Path,
    files: List[Dict[str, Any]],
    rename_map: Dict[str, str],
) -> None:
    local_dir.mkdir(parents=True, exist_ok=True)
    items = client.list_folder_all(folder_id)
    for item in items:
        safe_name = rename_map.get(item["id"], item["name"])
        local_path = local_dir / safe_name
        if item["type"] == "folder":
            _download_folder_recursive(client, item["id"], local_path, files, rename_map)
        else:
            if local_path.exists():
                print(f"  [skip] {safe_name} already exists locally")
            else:
                print(f"  downloading {safe_name} ...")
                client.download_file(item["id"], local_path)
            files.append({
                "drive_id": item["id"],
                "name": safe_name,
                "local_path": str(local_path),
                "sha256": sha256_file(local_path),
            })


def build_reference_index(
    files: List[Dict[str, Any]],
    local_root: Path,
) -> Dict[str, Optional[Path]]:
    """For each file, pick a reference: same basename elsewhere, or largest previously seen file."""
    by_basename: Dict[str, Path] = {}
    refs: Dict[str, Optional[Path]] = {}
    for f in files:
        local_path = Path(f["local_path"])
        basename = local_path.name
        if basename in by_basename and by_basename[basename] != local_path:
            refs[f["local_path"]] = by_basename[basename]
        else:
            # Use largest previously seen file as reference
            refs[f["local_path"]] = max(by_basename.values(), key=lambda p: p.stat().st_size) if by_basename else None
        by_basename[basename] = local_path
    return refs


def compress_folder(
    files: List[Dict[str, Any]],
    local_root: Path,
    archive_path: Path,
    threshold_q16: int = 32768,
    max_gccl_size: int = 100 * 1024 * 1024,
) -> Dict[str, Any]:
    """GCCL-compress downloaded files and pack into a .tar.zst archive."""
    refs = build_reference_index(files, local_root)
    manifest: Dict[str, Any] = {
        "schema": "mcp_drive_gccl_archive_v1",
        "created_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "base_dir": str(local_root),
        "files": [],
    }

    zstd_proc = subprocess.Popen(
        ["zstd", "-T0", "-19", "-q", "-f", "-o", str(archive_path)],
        stdin=subprocess.PIPE,
    )
    try:
        with tarfile.open(fileobj=zstd_proc.stdin, mode="w|") as tar:
            for f in files:
                src = Path(f["local_path"])
                rel = src.relative_to(local_root)
                ref = refs.get(f["local_path"])

                compressed_bytes, meta = gccl_compress_file(src, ref, threshold_q16, max_gccl_size)

                arc_name = f"{rel}.gccl" if meta["mode"] not in ("raw", "raw-oversize") else str(rel)
                if meta["mode"] in ("raw", "raw-oversize"):
                    tar.add(src, arcname=arc_name)
                else:
                    info = tarfile.TarInfo(name=arc_name)
                    info.size = len(compressed_bytes)
                    info.mtime = time.time()
                    tar.addfile(info, io.BytesIO(compressed_bytes))

                manifest["files"].append({
                    "path": str(rel),
                    "archive_path": arc_name,
                    "drive_id": f["drive_id"],
                    "sha256": f["sha256"],
                    **meta,
                })
    finally:
        if zstd_proc.stdin:
            zstd_proc.stdin.close()
        zstd_proc.wait()
        if zstd_proc.returncode != 0:
            raise RuntimeError(f"zstd exited with code {zstd_proc.returncode}")

    # Write manifest alongside
    manifest_path = archive_path.with_suffix(".manifest.json")
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")

    total_original = sum(m["original_size"] for m in manifest["files"])
    total_compressed = archive_path.stat().st_size + manifest_path.stat().st_size
    manifest["summary"] = {
        "file_count": len(files),
        "total_original_bytes": total_original,
        "archive_bytes": archive_path.stat().st_size,
        "manifest_bytes": manifest_path.stat().st_size,
        "total_output_bytes": total_compressed,
        "si_compression_ratio": total_original / total_compressed if total_compressed else 0,
    }
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    return manifest


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("folder_id", help="Google Drive folder ID to compress")
    parser.add_argument("--work-dir", default="/tmp/mcp-drive-gccl", help="Local scratch directory")
    parser.add_argument("--archive-name", help="Output archive name (default: folder_id.tar.zst)")
    parser.add_argument("--upload-parent", help="Drive folder ID to upload archive into")
    parser.add_argument("--trash-original", action="store_true", help="Move original folder to trash after upload")
    parser.add_argument("--threshold-q16", type=int, default=32768, help="GCCL threshold in Q16.16 (default 32768 = 0.5)")
    parser.add_argument("--max-gccl-size", type=int, default=100*1024*1024, help="Max bytes for GCCL compression; larger files stored raw (default 100 MiB)")
    parser.add_argument("--dry-run", action="store_true", help="Download and compress but do not upload or trash")
    args = parser.parse_args()

    archive_name = args.archive_name or f"{args.folder_id}.tar.zst"
    work_dir = Path(args.work_dir)
    local_clone = work_dir / "clone" / args.folder_id
    archive_path = work_dir / "out" / archive_name
    manifest_path = archive_path.with_suffix(".manifest.json")

    work_dir.mkdir(parents=True, exist_ok=True)
    local_clone.parent.mkdir(parents=True, exist_ok=True)
    archive_path.parent.mkdir(parents=True, exist_ok=True)

    env = {"GOOGLE_DRIVE_OAUTH_CREDENTIALS": "/home/allaun/gcp-oauth.keys.json"}
    client = DriveMcpClient(env=env)

    try:
        print(f"[*] Downloading folder {args.folder_id} to {local_clone}")
        files = download_folder_recursive(client, args.folder_id, local_clone)
        print(f"[+] Downloaded {len(files)} files")

        if not files:
            print("[-] No files found; aborting")
            return 1

        print(f"[*] GCCL-compressing into {archive_path}")
        manifest = compress_folder(files, local_clone, archive_path, args.threshold_q16, args.max_gccl_size)
        print(f"[+] Archive: {archive_path} ({archive_path.stat().st_size} bytes)")
        print(f"[+] Manifest: {manifest_path}")
        print(f"[+] SI compression ratio: {manifest['summary']['si_compression_ratio']:.2f}")

        if not args.dry_run:
            print("[*] Uploading archive to Google Drive")
            archive_upload_text = client.upload_file(archive_path, args.upload_parent)
            print(f"[+] Archive upload result: {archive_upload_text}")

            print("[*] Uploading manifest to Google Drive")
            manifest_upload_text = client.upload_file(manifest_path, args.upload_parent)
            print(f"[+] Manifest upload result: {manifest_upload_text}")

            if args.trash_original:
                print(f"[*] Trashing original folder {args.folder_id}")
                client.delete_item(args.folder_id)
                print("[+] Original trashed")
        else:
            print("[DRY RUN] No upload or trash performed")

        # Print manifest summary
        print(json.dumps(manifest["summary"], indent=2))
        return 0

    finally:
        client.close()


if __name__ == "__main__":
    sys.exit(main())