import os
import shutil
import zipfile
import json
from pathlib import Path
from datetime import datetime
from typing import List, Optional, Dict, Any
import importlib.util


class ModuleBackup:
    """
    Backup and restore manager for Python modules.

    - Give only the module name (string).
    - Automatically locates the module path using importlib.
    - Creates a hidden backup folder `.backup_<modulename>` near the module.
    - Backups are timestamped automatically.
    """

    def __init__(self, module_name: str):
        self.module_name = module_name

        spec = importlib.util.find_spec(module_name)
        if spec is None or spec.origin is None:
            raise ModuleNotFoundError(f"Cannot find module: {module_name}")

        self.module_path = Path(spec.origin).resolve()
        if not self.module_path.exists():
            raise FileNotFoundError(f"Module not found at {self.module_path}")

        # If it's a package (__init__.py), use the parent folder
        if self.module_path.name == "__init__.py":
            self.module_path = self.module_path.parent

        # Default backup folder: .backup_<module>
        self.backups_root = self.module_path.parent / f".backup_{self.module_name}"
        self.backups_root.mkdir(parents=True, exist_ok=True)

        self.index_file = self.backups_root / "backups_index.json"
        if not self.index_file.exists():
            self._write_index([])

    # ---------- internals ----------
    def _read_index(self) -> List[Dict[str, Any]]:
        try:
            with open(self.index_file, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return []

    def _write_index(self, data: List[Dict[str, Any]]) -> None:
        with open(self.index_file, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    def _make_stamp(self) -> str:
        return datetime.now().strftime("%Y%m%dT%H%M%SZ")

    def _backup_name(self, stamp: str) -> str:
        return f"{self.module_name}_{stamp}"

    # ---------- user-facing ----------
    def backup(
        self, compress: bool = True, message: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Create a new backup for the module.
        - compress: If True, saves as .zip archive, else saves as folder.
        - message: Optional note for metadata.
        Returns a dict with backup info.
        """
        stamp = self._make_stamp()
        base_name = self._backup_name(stamp)
        target_folder = self.backups_root / base_name

        if target_folder.exists():
            raise FileExistsError(f"Backup target already exists: {target_folder}")

        # Copy file or folder
        if self.module_path.is_file():
            target_folder.mkdir(parents=True, exist_ok=False)
            dest = target_folder / self.module_path.name
            shutil.copy2(self.module_path, dest)
        else:
            shutil.copytree(self.module_path, target_folder)

        backup_entry = {
            "stamp": stamp,
            "name": base_name,
            "created_at": datetime.now().isoformat() + "Z",
            "original": str(self.module_path),
            "compress": bool(compress),
            "message": message or "",
        }

        # Compress if requested
        archive_path = None
        if compress:
            archive_path = self.backups_root / f"{base_name}.zip"
            with zipfile.ZipFile(
                archive_path, "w", compression=zipfile.ZIP_DEFLATED
            ) as zf:
                for p in target_folder.rglob("*"):
                    rel = p.relative_to(target_folder)
                    zf.write(p, arcname=str(Path(base_name) / rel))
            # Remove folder after compression
            shutil.rmtree(target_folder)
            backup_entry["archive"] = str(archive_path)
        else:
            backup_entry["archive"] = None

        index = self._read_index()
        index.insert(0, backup_entry)
        self._write_index(index)

        return {"status": "ok", "backup": backup_entry}

    def list_backups(self) -> List[Dict[str, Any]]:
        """Return a list of available backups (newest first)."""
        return self._read_index()

    def restore(
        self, stamp: Optional[str] = None, overwrite: bool = True
    ) -> Dict[str, Any]:
        """
        Restore a backup.
        - stamp: If given, restore that specific backup; otherwise restores the newest.
        - overwrite: If True, deletes the original before restoring.
        """
        index = self._read_index()
        if not index:
            raise FileNotFoundError("No backups available to restore.")

        entry = None
        if stamp:
            entry = next((e for e in index if e["stamp"] == stamp), None)
            if entry is None:
                raise FileNotFoundError(f"No backup found with stamp '{stamp}'")
        else:
            entry = index[0]

        if entry.get("archive"):
            archive_path = Path(entry["archive"])
            if not archive_path.exists():
                raise FileNotFoundError(f"Archive not found: {archive_path}")
            tmp_extract = self.backups_root / f"__extract__{entry['stamp']}"
            if tmp_extract.exists():
                shutil.rmtree(tmp_extract)
            tmp_extract.mkdir()
            with zipfile.ZipFile(archive_path, "r") as zf:
                zf.extractall(tmp_extract)
            src_root = tmp_extract / entry["name"]
        else:
            src_root = self.backups_root / entry["name"]
            if not src_root.exists():
                raise FileNotFoundError(f"Backup folder missing: {src_root}")

        # Overwrite original if requested
        if overwrite and self.module_path.exists():
            if self.module_path.is_file():
                self.module_path.unlink()
            else:
                shutil.rmtree(self.module_path)

        # Copy back
        if src_root.is_dir():
            shutil.copytree(src_root, self.module_path)
        else:
            self.module_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src_root, self.module_path)

        if entry.get("archive"):
            shutil.rmtree(tmp_extract)

        return {
            "status": "restored",
            "from": entry["stamp"],
            "original": entry["original"],
        }
