#!/usr/bin/env python3
"""SophClaw state restore tool.

Extracts a SophClaw backup ZIP and merges its contents into the current
state directory (~/.openclaw by default).

Usage:
    python3 restore.py [backup.zip]

If no path is given, auto-discovers the latest ZIP in ~/.openclaw/media/inbound/.
"""

import json
import os
import shutil
import sys
import tempfile
import zipfile
from pathlib import Path


def resolve_state_dir() -> Path:
    """Resolve the SophClaw state directory, same logic as SophClaw's resolveStateDir."""
    env_override = os.environ.get("OPENCLAW_STATE_DIR")
    if env_override:
        return Path(env_override).expanduser().resolve()

    home = Path.home()
    candidates = [
        home / ".openclaw",
        home / ".clawdbot",
        home / ".moltbot",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate

    return home / ".openclaw"


def auto_discover_zip(state_dir: Path) -> Path | None:
    """Find the most recent ZIP file in the media/inbound directory."""
    inbound = state_dir / "media" / "inbound"
    if not inbound.is_dir():
        return None

    zip_files = sorted(
        inbound.glob("*.zip"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    if not zip_files:
        return None
    return zip_files[0]


# Config keys to RESTORE from the backup (user-created content and preferences).
RESTORE_BACKUP_KEYS = {
    "agents",
    "models",
    "skills",
    "tools",
    "commands",
    "messages",
    "session",
    "identity",
    "cron",
}


def deep_merge(base: dict, overlay: dict) -> dict:
    """Recursively merge overlay into base. overlay values win on conflict."""
    result = dict(base)
    for key, value in overlay.items():
        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
            result[key] = deep_merge(result[key], value)
        else:
            result[key] = value
    return result


def merge_openclaw_json(backup_path: Path, state_dir: Path) -> dict:
    """Merge backup openclaw.json with the current one.

    Backup values take precedence for user-config keys (agents, models, etc.).
    Container runtime keys (gateway, channels, plugins) are kept from current.
    """
    current_path = state_dir / "openclaw.json"
    if not current_path.exists():
        with open(backup_path, "r", encoding="utf-8") as f:
            return json.load(f)

    with open(current_path, "r", encoding="utf-8") as f:
        current = json.load(f)
    with open(backup_path, "r", encoding="utf-8") as f:
        backup = json.load(f)

    merged = dict(current)

    for key in RESTORE_BACKUP_KEYS:
        if key in backup:
            if key in merged and isinstance(merged[key], dict) and isinstance(backup[key], dict):
                merged[key] = deep_merge(merged[key], backup[key])
            else:
                merged[key] = backup[key]

    return merged


def merge_sessions_json(backup_path: Path, current_path: Path) -> tuple[int, int]:
    """Merge backup sessions.json into the current one.

    Backup entries win on session ID conflict.
    Returns (backup_session_count, merged_total_count).
    """
    current_sessions: dict = {}
    if current_path.exists():
        with open(current_path, "r", encoding="utf-8") as f:
            current_sessions = json.load(f)

    with open(backup_path, "r", encoding="utf-8") as f:
        backup_sessions = json.load(f)

    merged = dict(current_sessions)
    for session_id, entry in backup_sessions.items():
        merged[session_id] = entry

    current_path.parent.mkdir(parents=True, exist_ok=True)
    with open(current_path, "w", encoding="utf-8") as f:
        json.dump(merged, f, ensure_ascii=False, indent=2)
        f.write("\n")

    return len(backup_sessions), len(merged)


def copy_file(src: Path, dst: Path) -> bool:
    """Copy a file from src to dst, creating parent dirs as needed."""
    try:
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)
        return True
    except OSError as e:
        print(f"  [WARN] Failed to copy {src} -> {dst}: {e}", file=sys.stderr)
        return False


def copy_directory(src: Path, dst: Path) -> tuple[int, int]:
    """Recursively copy directory contents from src to dst."""
    copied = 0
    skipped = 0
    dst.mkdir(parents=True, exist_ok=True)
    for item in src.rglob("*"):
        rel = item.relative_to(src)
        target = dst / rel
        if item.is_dir():
            target.mkdir(parents=True, exist_ok=True)
            continue
        if copy_file(item, target):
            copied += 1
        else:
            skipped += 1
    return copied, skipped


def resolve_backup_path(args: list[str], state_dir: Path) -> tuple[Path, str]:
    """Resolve the backup ZIP path from CLI args or auto-discovery.

    Returns (path, source_description).
    """
    if len(args) >= 2:
        user_path = Path(args[1]).resolve()
        if not user_path.exists():
            print(f"Error: backup file not found: {user_path}", file=sys.stderr)
            sys.exit(1)
        if not zipfile.is_zipfile(user_path):
            print(f"Error: not a valid ZIP file: {user_path}", file=sys.stderr)
            sys.exit(1)
        return user_path, "user-specified"

    # Auto-discover from media/inbound
    discovered = auto_discover_zip(state_dir)
    if discovered is None:
        print(
            "Error: no backup ZIP path provided and no ZIP files found in "
            f"{state_dir / 'media' / 'inbound'}",
            file=sys.stderr,
        )
        print(
            "Usage: python3 restore.py <backup.zip>\n"
            "       or upload a ZIP via SophClaw chat and run without arguments.",
            file=sys.stderr,
        )
        sys.exit(1)

    return discovered, "auto-discovered"


def main():
    state_dir = resolve_state_dir()
    state_dir.mkdir(parents=True, exist_ok=True)

    zip_path, source = resolve_backup_path(sys.argv, state_dir)

    print(f"State directory: {state_dir}")
    print(f"Backup ZIP ({source}): {zip_path}")
    print(f"  Size: {zip_path.stat().st_size / 1024:.1f} KB")
    print()

    summary = {
        "agents_restored": [],
        "sessions_total": 0,
        "sessions_merged": 0,
        "files_copied": 0,
        "files_skipped": 0,
        "workspaces_restored": [],
    }

    with tempfile.TemporaryDirectory(prefix="sophclaw-restore-") as tmpdir:
        tmp = Path(tmpdir)
        print("Extracting ZIP...")
        with zipfile.ZipFile(zip_path, "r") as zf:
            zf.extractall(tmp)

        top_items = list(tmp.iterdir())
        print(f"  Extracted {len(top_items)} top-level items")

        # === 1. openclaw.json ===
        backup_config = tmp / "openclaw.json"
        if backup_config.exists():
            print("\n--- Config ---")
            print("  Merging openclaw.json...")
            merged = merge_openclaw_json(backup_config, state_dir)

            config_target = state_dir / "openclaw.json"
            with open(config_target, "w", encoding="utf-8") as f:
                json.dump(merged, f, ensure_ascii=False, indent=2)
                f.write("\n")
            print(f"  Written: {config_target}")

            agents_list = merged.get("agents", {}).get("list", [])
            for agent in agents_list:
                summary["agents_restored"].append(agent.get("id", "unknown"))
            print(f"  Agents restored: {', '.join(summary['agents_restored']) or 'none'}")
        else:
            print("  [WARN] No openclaw.json found in backup")

        # === 2. Agents (sessions, models, runtime-stats) ===
        backup_agents = tmp / "agents"
        if backup_agents.is_dir():
            print("\n--- Agents ---")
            current_agents = state_dir / "agents"

            for agent_dir in sorted(backup_agents.iterdir()):
                if not agent_dir.is_dir():
                    continue
                agent_name = agent_dir.name

                # Merge sessions.json
                backup_sessions_json = agent_dir / "sessions" / "sessions.json"
                current_sessions_json = current_agents / agent_name / "sessions" / "sessions.json"
                if backup_sessions_json.exists():
                    backup_count, merged_count = merge_sessions_json(
                        backup_sessions_json, current_sessions_json
                    )
                    summary["sessions_merged"] += backup_count
                    summary["sessions_total"] += merged_count
                    print(
                        f"  [{agent_name}] sessions: {backup_count} from backup, "
                        f"{merged_count} total after merge"
                    )

                # Copy session transcript .jsonl files
                backup_sessions_dir = agent_dir / "sessions"
                if backup_sessions_dir.is_dir():
                    jsonl_count = 0
                    for session_file in backup_sessions_dir.iterdir():
                        if session_file.suffix == ".jsonl" and session_file.name != "sessions.json":
                            target = current_agents / agent_name / "sessions" / session_file.name
                            if copy_file(session_file, target):
                                jsonl_count += 1
                    if jsonl_count > 0:
                        print(f"  [{agent_name}] transcripts: {jsonl_count} files copied")

                # Copy non-session agent files (models.json, runtime-stats, etc.)
                for item in agent_dir.rglob("*"):
                    if item.is_dir():
                        continue
                    rel = item.relative_to(agent_dir)
                    if rel.parts[0] == "sessions":
                        continue
                    target = current_agents / agent_name / rel
                    if copy_file(item, target):
                        summary["files_copied"] += 1
        else:
            print("\n  [INFO] No agents directory in backup")

        # === 3. Workspace directories ===
        print("\n--- Workspaces ---")
        for item in sorted(tmp.iterdir()):
            if not item.is_dir():
                continue
            if item.name.startswith("workspace"):
                ws_name = item.name
                current_ws = state_dir / ws_name
                copied, skipped = copy_directory(item, current_ws)
                summary["files_copied"] += copied
                summary["files_skipped"] += skipped
                summary["workspaces_restored"].append(ws_name)
                print(f"  [{ws_name}] {copied} files, {skipped} skipped")

        # === 4. Memory SQLite files ===
        backup_memory = tmp / "memory"
        if backup_memory.is_dir():
            print("\n--- Memory ---")
            current_memory = state_dir / "memory"
            for db_file in sorted(backup_memory.iterdir()):
                if db_file.suffix == ".sqlite":
                    target = current_memory / db_file.name
                    if copy_file(db_file, target):
                        summary["files_copied"] += 1
                        print(f"  [{db_file.name}] restored")

        # === 5. Loose files (identity, canvas, logs, media, config backups, etc.) ===
        LOOSE_PATTERNS = [
            "identity",
            "canvas",
            "logs",
            "media",
            "openclaw.json.bak",
            "openclaw.json.bak.1",
            "openclaw.json.bak.2",
            "openclaw.json.bak.3",
            "openclaw.json.bak.4",
            "openclaw.json.last-good",
            "jwt.json",
            "update-check.json",
        ]

        print("\n--- Other Files ---")
        for pattern in LOOSE_PATTERNS:
            src = tmp / pattern
            if src.exists():
                dst = state_dir / pattern
                if src.is_dir():
                    copied, skipped = copy_directory(src, dst)
                    summary["files_copied"] += copied
                    summary["files_skipped"] += skipped
                    print(f"  [{pattern}/] {copied} files")
                else:
                    if copy_file(src, dst):
                        summary["files_copied"] += 1
                        print(f"  [{pattern}] copied")

        # Handle any remaining items not explicitly covered
        handled = set(LOOSE_PATTERNS) | {"agents", "memory", "openclaw.json"}
        for ws_name in summary["workspaces_restored"]:
            handled.add(ws_name)

        for item in sorted(tmp.iterdir()):
            if item.name in handled:
                continue
            dst = state_dir / item.name
            if item.is_dir():
                copied, skipped = copy_directory(item, dst)
                summary["files_copied"] += copied
                summary["files_skipped"] += skipped
                if copied > 0:
                    print(f"  [{item.name}/] {copied} files")
            else:
                if copy_file(item, dst):
                    summary["files_copied"] += 1
                    print(f"  [{item.name}] copied")

    # === Summary ===
    print()
    print("=" * 50)
    print("Restore Complete")
    print("=" * 50)
    print(f"  Agents:      {', '.join(summary['agents_restored']) or 'none'}")
    print(f"  Sessions:    {summary['sessions_merged']} from backup ({summary['sessions_total']} total)")
    print(f"  Workspaces:  {', '.join(summary['workspaces_restored']) or 'none'}")
    print(f"  Files:       {summary['files_copied']} copied, {summary['files_skipped']} skipped")
    print()

    # Touch config to trigger chokidar-based hot reload
    config_target = state_dir / "openclaw.json"
    if config_target.exists():
        os.utime(config_target, None)
        print("Config updated. SophClaw will auto-reload within a few seconds.")
        print("Refresh your browser page to see the restored state.")

    return 0


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