import json
import re
import shutil
import subprocess
from typing import List, Optional

from sam.managers import Manager
from sam.actions import Action


class SystemManager(Manager):
    NVIDIA_IMAGE = "quay.io/stillhq/stillos-workstation-nvidia"
    title: str
    manager_id: str
    current_action: Action = None

    @property
    def manager_type(self) -> str:
        return "system"

    def __init__(self, title: str = "System", manager_id: str = "bootc"):
        super().__init__(title, manager_id)

    def _run_bootc(self, args: List[str]) -> subprocess.CompletedProcess:
        """Run a bootc command and return the result."""
        cmd = ["bootc"] + args
        return subprocess.run(cmd, capture_output=True, text=True)

    def _run_rpm_ostree(self, args: List[str]) -> subprocess.CompletedProcess:
        """Run an rpm-ostree command and return the result."""
        cmd = ["rpm-ostree"] + args
        return subprocess.run(cmd, capture_output=True, text=True)

    def _get_bootc_status(self) -> Optional[dict]:
        """Get bootc status as JSON."""
        result = self._run_bootc(["status", "--json"])
        if result.returncode != 0:
            return None
        try:
            return json.loads(result.stdout)
        except json.JSONDecodeError:
            return None

    def _get_rpm_ostree_status(self) -> Optional[dict]:
        """Get rpm-ostree status as JSON."""
        if shutil.which("rpm-ostree") is None:
            return None

        result = self._run_rpm_ostree(["status", "--json"])
        if result.returncode != 0:
            return None
        try:
            return json.loads(result.stdout)
        except json.JSONDecodeError:
            return None

    def _get_booted_rpm_ostree_deployment(self) -> Optional[dict]:
        status = self._get_rpm_ostree_status()
        if status is None:
            return None

        for deployment in status.get("deployments", []):
            if deployment.get("booted"):
                return deployment
        return None

    def _has_rpm_ostree_overlay(self) -> bool:
        """Return True when the booted deployment has rpm-ostree overlays."""
        deployment = self._get_booted_rpm_ostree_deployment()
        if deployment is None:
            return False

        overlay_keys = (
            "requested-packages",
            "requested-local-packages",
            "requested-base-removals",
            "requested-base-local-replacements",
            "base-local-replacements",
        )
        return any(len(deployment.get(key, [])) > 0 for key in overlay_keys)

    def _has_update_available(self) -> bool:
        """Check if a system update is available."""
        if self._has_rpm_ostree_overlay():
            result = self._run_rpm_ostree(["upgrade", "--check"])
            return result.returncode != 0

        # bootc upgrade --check returns non-zero exit code if an update is available
        result = self._run_bootc(["upgrade", "--check"])
        return result.returncode != 0

    def install(self, action: Action):
        """Run allowlisted system installation transactions."""
        self.current_action = action
        if action.package_id != "nvidia-driver" or set(action.tags) != {"system", "nvidia"}:
            action.error = "Unsupported system installation transaction"
            return

        action.status = "Switching to the NVIDIA workstation image..."
        action.progress = 0
        result = self._run_bootc(["switch", self.NVIDIA_IMAGE])
        if result.returncode != 0:
            action.error = result.stderr or result.stdout or "Failed to switch to the NVIDIA workstation image"
            return

        action.progress = 100
        action.status = "NVIDIA drivers staged. Reboot to apply."

    def remove(self, action: Action):
        """Remove is a no-op for system manager."""
        self.current_action = action
        action.progress = 100
        # No-op - cannot remove the system

    def _parse_layer_progress(self, output: str) -> Optional[tuple[int, str]]:
        """Parse layer progress from output, returning (percentage, "done/total") or None.
        
        Looks for "Fetching layers" line with format like "X/Y" at the end.
        """
        # Split by lines (handle both \n and \r for in-place updates)
        lines = re.split(r'[\r\n]+', output)
        
        for line in lines:
            # Look for the fetching layers line with progress like "0/48" or "23/48"
            if "Fetching layers" in line or "layers" in line.lower():
                match = re.search(r'(\d+)/(\d+)', line)
                if match:
                    done = int(match.group(1))
                    total = int(match.group(2))
                    if total > 0:
                        percentage = int((done / total) * 100)
                        progress_str = f"{done}/{total}"
                        return (percentage, progress_str)
        return None

    def _upgrade_command(self) -> tuple[list[str], str]:
        if self._has_rpm_ostree_overlay():
            return ["rpm-ostree", "upgrade"], "rpm-ostree"
        return ["bootc", "upgrade"], "bootc"

    def update(self, action: Action):
        """Update the system using bootc or rpm-ostree when overlays require it."""
        self.current_action = action
        cmd, backend = self._upgrade_command()
        action.status = f"Upgrading system with {backend}..."
        action.progress = 0

        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1  # Line buffered
        )

        output_buffer = ""

        # Read output character by character to handle in-place updates
        while True:
            char = process.stdout.read(1)
            if not char:
                if process.poll() is not None:
                    break
                continue

            output_buffer += char

            # Parse progress whenever we have enough output
            if char in ('\n', '\r') or len(output_buffer) > 100:
                result = self._parse_layer_progress(output_buffer)
                if result is not None:
                    percentage, progress_str = result
                    action.progress = min(percentage, 99)  # Cap at 99 until complete
                    action.status = f"Fetching layers {progress_str}"
                # Keep last 500 chars to handle multi-line parsing
                if len(output_buffer) > 500:
                    output_buffer = output_buffer[-500:]

        # Final check for progress
        result = self._parse_layer_progress(output_buffer)
        if result is None:
            action.progress = 99

        # Wait for process to complete
        process.wait()

        if process.returncode != 0:
            stderr = process.stderr.read()
            action.error = stderr or "Failed to upgrade system"
        else:
            action.progress = 100
            action.status = "System upgrade staged. Reboot to apply."

    def check_update(self, package: str) -> bool:
        """Check if a specific package has an update (only 'system' is valid)."""
        if package != "system":
            return False
        return self._has_update_available()

    def check_updates(self) -> List[str]:
        """Return ['system'] if a system update is available, empty list otherwise."""
        if self._has_update_available():
            return ["system"]
        return []

    def check_installed(self) -> List[str]:
        """Return ['system'] as the system is always 'installed'."""
        # Verify bootc or rpm-ostree is actually managing this system
        if self._get_bootc_status() is not None or self._get_rpm_ostree_status() is not None:
            return ["system"]
        return []

    def bare_app_info(self, package: str) -> Optional[dict]:
        """Return basic info for the system package."""
        if package != "system":
            return None

        status = self._get_bootc_status()
        booted_image = "Unknown"

        if status is not None:
            booted = status.get("status", {}).get("booted", {})
            image = booted.get("image", {})
            booted_image = image.get("image", {}).get("image", "Unknown")
        else:
            deployment = self._get_booted_rpm_ostree_deployment()
            if deployment is not None:
                booted_image = deployment.get("origin", "rpm-ostree deployment")

        backend = "rpm-ostree" if self._has_rpm_ostree_overlay() else "bootc"

        return {
            "app_id": "system",
            "name": "System",
            "author": "System",
            "primary_src": self.manager_id,
            "src_package_name": "system",
            "icon_url": "",
            "summary": f"{backend} managed system image: {booted_image}",
            "description": f"The base operating system image managed by {backend}.",
        }                   
