"""Core update manager — download, verify, apply, rollback."""
from __future__ import annotations

import hashlib
import json
import logging
import os
import shutil
import subprocess
import sys
import tarfile
import tempfile
import time
from datetime import datetime
from pathlib import Path
from typing import Optional

from mesh.updates.models import (
    UpdateChannel, UpdateHistory, UpdateHistoryEntry, UpdateInfo, UpdateStatus,
)

logger = logging.getLogger(__name__)

# Where the mesh package lives (parent of the mesh/ directory)
_DEFAULT_INSTALL_DIR = Path(__file__).resolve().parents[2]


class UpdateManager:
    """Handles downloading, verifying, applying, and rolling back updates."""

    def __init__(
        self,
        current_version: str = "0.1.0",
        install_dir: Optional[Path] = None,
        data_dir: Optional[Path] = None,
        update_server_url: str = "",
    ):
        self.current_version = current_version
        self.install_dir = Path(install_dir) if install_dir else _DEFAULT_INSTALL_DIR
        self.data_dir = Path(data_dir) if data_dir else self.install_dir / "data"
        self.update_server_url = update_server_url.rstrip("/")

        self.backup_dir = self.data_dir / "backups"
        self.backup_dir.mkdir(parents=True, exist_ok=True)

        self.history_path = self.data_dir / "update_history.json"
        self._history = self._load_history()

    # ── History persistence ──────────────────────────────────────

    def _load_history(self) -> UpdateHistory:
        if self.history_path.exists():
            try:
                raw = json.loads(self.history_path.read_text())
                return UpdateHistory(**raw)
            except Exception as e:
                logger.warning("Failed to load update history: %s", e)
        return UpdateHistory(current_version=self.current_version)

    def _save_history(self) -> None:
        self.history_path.parent.mkdir(parents=True, exist_ok=True)
        self.history_path.write_text(
            self._history.model_dump_json(indent=2)
        )

    @property
    def history(self) -> UpdateHistory:
        return self._history

    # ── Check for updates ────────────────────────────────────────

    def check_for_updates(
        self, channel: str = "stable"
    ) -> Optional[UpdateInfo]:
        """Check the update server for a newer version.

        Returns UpdateInfo if an update is available, None otherwise.
        Works without a server — returns None (manual-only mode).
        """
        if not self.update_server_url:
            logger.debug("No update server configured — manual mode only")
            return None

        import urllib.request
        import urllib.error

        url = f"{self.update_server_url}/updates/latest?channel={channel}"
        try:
            req = urllib.request.Request(url, headers={"Accept": "application/json"})
            with urllib.request.urlopen(req, timeout=10) as resp:
                data = json.loads(resp.read())
                info = UpdateInfo(**data)
                if self._is_newer(info.version):
                    logger.info("Update available: %s -> %s", self.current_version, info.version)
                    return info
                logger.debug("Already on latest version %s", self.current_version)
                return None
        except urllib.error.URLError as e:
            logger.warning("Failed to check for updates: %s", e)
            return None
        except Exception as e:
            logger.warning("Update check error: %s", e)
            return None

    def _is_newer(self, version: str) -> bool:
        """Compare semantic versions. Returns True if version > current."""
        try:
            current = tuple(int(x) for x in self.current_version.split("."))
            new = tuple(int(x) for x in version.split("."))
            return new > current
        except (ValueError, AttributeError):
            return False

    # ── Download ─────────────────────────────────────────────────

    def download_update(self, update_info: UpdateInfo) -> Path:
        """Download an update package to a temp directory.

        Returns the path to the downloaded tarball.
        """
        if not update_info.download_url:
            raise ValueError("No download URL in update info")

        import urllib.request

        tmp_dir = Path(tempfile.mkdtemp(prefix="mesh-update-"))
        dest = tmp_dir / f"mesh-{update_info.version}.tar.gz"

        logger.info("Downloading update %s from %s", update_info.version, update_info.download_url)
        urllib.request.urlretrieve(update_info.download_url, str(dest))

        actual_size = dest.stat().st_size
        if update_info.size_bytes and actual_size != update_info.size_bytes:
            dest.unlink()
            raise ValueError(
                f"Size mismatch: expected {update_info.size_bytes}, got {actual_size}"
            )

        logger.info("Downloaded %s (%.1f MB)", dest.name, actual_size / 1048576)
        return dest

    # ── Verify ───────────────────────────────────────────────────

    @staticmethod
    def verify_update(package_path: Path, expected_sha256: str) -> bool:
        """Verify package integrity via SHA-256."""
        if not expected_sha256:
            logger.warning("No SHA-256 provided — skipping verification")
            return True

        sha256 = hashlib.sha256()
        with open(package_path, "rb") as f:
            for chunk in iter(lambda: f.read(65536), b""):
                sha256.update(chunk)
        actual = sha256.hexdigest()
        ok = actual == expected_sha256
        if not ok:
            logger.error("SHA-256 mismatch: expected %s, got %s", expected_sha256, actual)
        return ok

    # ── Apply ────────────────────────────────────────────────────

    def apply_update(self, package_path: Path, restart: bool = True) -> UpdateStatus:
        """Apply an update package.

        1. Create backup of current installation
        2. Extract package over install dir
        3. Optionally restart the service

        Returns the final status.
        """
        package_path = Path(package_path)
        if not package_path.exists():
            raise FileNotFoundError(f"Package not found: {package_path}")

        # Determine version from tarball name or metadata inside
        version = self._extract_version(package_path)

        entry = UpdateHistoryEntry(
            version=version,
            status=UpdateStatus.APPLYING,
            applied_at=datetime.utcnow(),
        )

        try:
            # Step 1: Backup
            backup_path = self._create_backup()
            entry.backup_path = str(backup_path)
            logger.info("Backup created at %s", backup_path)

            # Step 2: Extract
            self._extract_package(package_path)
            logger.info("Package extracted to %s", self.install_dir)

            # Step 3: Update version tracking
            entry.status = UpdateStatus.COMPLETE
            self.current_version = version
            self._history.current_version = version
            self._history.entries.append(entry)
            self._save_history()

            logger.info("Update to %s applied successfully", version)

            # Step 4: Restart if requested
            if restart:
                self._schedule_restart()

            return UpdateStatus.COMPLETE

        except Exception as e:
            logger.error("Update failed: %s", e)
            entry.status = UpdateStatus.FAILED
            entry.error = str(e)
            self._history.entries.append(entry)
            self._save_history()

            # Auto-rollback on failure
            if entry.backup_path:
                try:
                    self._restore_backup(Path(entry.backup_path))
                    entry.status = UpdateStatus.ROLLED_BACK
                    entry.rolled_back_at = datetime.utcnow()
                    self._save_history()
                    logger.info("Auto-rollback completed")
                except Exception as rb_err:
                    logger.error("Auto-rollback also failed: %s", rb_err)

            return entry.status

    def _extract_version(self, package_path: Path) -> str:
        """Extract version string from package."""
        name = package_path.stem  # e.g. mesh-0.2.0.tar
        if name.endswith(".tar"):
            name = name[:-4]
        parts = name.split("-", 1)
        if len(parts) == 2 and parts[0] == "mesh":
            return parts[1]

        # Try reading from metadata inside the tarball
        try:
            with tarfile.open(package_path, "r:*") as tf:
                for member in tf.getnames():
                    if member.endswith("VERSION") or member.endswith("version.txt"):
                        f = tf.extractfile(member)
                        if f:
                            return f.read().decode().strip()
        except Exception:
            pass

        return "unknown"

    def _create_backup(self) -> Path:
        """Back up the current mesh/ directory."""
        ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        backup_name = f"mesh-{self.current_version}-{ts}"
        backup_path = self.backup_dir / backup_name

        mesh_src = self.install_dir / "mesh"
        if mesh_src.is_dir():
            shutil.copytree(mesh_src, backup_path / "mesh", symlinks=True)
        else:
            backup_path.mkdir(parents=True, exist_ok=True)

        # Also back up pyproject.toml if it exists
        pyproject = self.install_dir / "pyproject.toml"
        if pyproject.exists():
            backup_path.mkdir(parents=True, exist_ok=True)
            shutil.copy2(pyproject, backup_path / "pyproject.toml")

        # Prune old backups (keep most recent N)
        self._prune_backups(keep=3)

        return backup_path

    def _prune_backups(self, keep: int = 3) -> None:
        """Remove old backups, keeping the most recent `keep`."""
        if not self.backup_dir.exists():
            return
        backups = sorted(
            [d for d in self.backup_dir.iterdir() if d.is_dir()],
            key=lambda d: d.stat().st_mtime,
            reverse=True,
        )
        for old in backups[keep:]:
            logger.info("Pruning old backup: %s", old.name)
            shutil.rmtree(old, ignore_errors=True)

    def _extract_package(self, package_path: Path) -> None:
        """Extract the update tarball over the install directory."""
        with tarfile.open(package_path, "r:*") as tf:
            # Security: check for path traversal
            for member in tf.getmembers():
                member_path = Path(self.install_dir) / member.name
                try:
                    member_path.resolve().relative_to(self.install_dir.resolve())
                except ValueError:
                    raise ValueError(
                        f"Path traversal detected in archive: {member.name}"
                    )
            tf.extractall(self.install_dir, filter="data")

    def _restore_backup(self, backup_path: Path) -> None:
        """Restore a backup by copying files back."""
        mesh_backup = backup_path / "mesh"
        mesh_dst = self.install_dir / "mesh"

        if mesh_backup.is_dir():
            if mesh_dst.exists():
                shutil.rmtree(mesh_dst)
            shutil.copytree(mesh_backup, mesh_dst, symlinks=True)

        pyproject_backup = backup_path / "pyproject.toml"
        if pyproject_backup.exists():
            shutil.copy2(pyproject_backup, self.install_dir / "pyproject.toml")

    def _schedule_restart(self) -> None:
        """Restart the current process after a short delay.

        Spawns a detached subprocess that waits 2 seconds then re-execs
        the current process with the same arguments.
        """
        logger.info("Scheduling restart in 2 seconds")
        restart_script = (
            f"import time, os, sys; time.sleep(2); "
            f"os.execv(sys.executable, {[sys.executable] + sys.argv!r})"
        )
        subprocess.Popen(
            [sys.executable, "-c", restart_script],
            start_new_session=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )

    # ── Rollback ─────────────────────────────────────────────────

    def rollback(self) -> UpdateStatus:
        """Rollback to the most recent backup.

        Returns the resulting status.
        """
        # Find the last successful or failed entry that has a backup
        for entry in reversed(self._history.entries):
            if entry.backup_path and Path(entry.backup_path).exists():
                logger.info("Rolling back to backup: %s", entry.backup_path)
                try:
                    self._restore_backup(Path(entry.backup_path))
                    entry.status = UpdateStatus.ROLLED_BACK
                    entry.rolled_back_at = datetime.utcnow()

                    # Restore version
                    # Parse version from backup dir name: mesh-X.Y.Z-timestamp
                    backup_name = Path(entry.backup_path).name
                    parts = backup_name.split("-")
                    if len(parts) >= 2:
                        self.current_version = parts[1]
                        self._history.current_version = self.current_version

                    self._save_history()
                    logger.info("Rollback to %s complete", self.current_version)
                    return UpdateStatus.ROLLED_BACK
                except Exception as e:
                    logger.error("Rollback failed: %s", e)
                    return UpdateStatus.FAILED

        logger.error("No backup available for rollback")
        return UpdateStatus.FAILED

    # ── Utility ──────────────────────────────────────────────────

    @staticmethod
    def compute_sha256(path: Path) -> str:
        """Compute SHA-256 of a file."""
        sha256 = hashlib.sha256()
        with open(path, "rb") as f:
            for chunk in iter(lambda: f.read(65536), b""):
                sha256.update(chunk)
        return sha256.hexdigest()
