"""Mags client for interacting with the Magpie VM infrastructure API."""

from __future__ import annotations

import os
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import requests


class MagsError(Exception):
    """Raised when the Mags API returns an error."""

    def __init__(self, message: str, status_code: int | None = None):
        super().__init__(message)
        self.status_code = status_code


class Mags:
    """Client for the Mags API.

    Args:
        api_token: API token. Falls back to ``MAGS_API_TOKEN`` or ``MAGS_TOKEN`` env vars.
        api_url: API base URL. Falls back to ``MAGS_API_URL`` env var or
            ``https://api.magpiecloud.com``.
        timeout: Default request timeout in seconds.
    """

    def __init__(
        self,
        api_token: str | None = None,
        api_url: str | None = None,
        timeout: int = 30,
    ):
        self.api_url = (
            api_url
            or os.environ.get("MAGS_API_URL")
            or "https://api.magpiecloud.com"
        ).rstrip("/")

        self.api_token = (
            api_token
            or os.environ.get("MAGS_API_TOKEN")
            or os.environ.get("MAGS_TOKEN")
        )
        if not self.api_token:
            raise MagsError(
                "API token required. Set MAGS_API_TOKEN env var or pass api_token."
            )

        self.timeout = timeout
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {self.api_token}",
                "Content-Type": "application/json",
            }
        )

    # ── helpers ──────────────────────────────────────────────────────

    def _request(
        self,
        method: str,
        path: str,
        json: Any = None,
        params: dict | None = None,
        timeout: int | None = None,
    ) -> Any:
        url = f"{self.api_url}/api/v1{path}"
        resp = self._session.request(
            method,
            url,
            json=json,
            params=params,
            timeout=timeout or self.timeout,
        )
        if resp.status_code >= 400:
            try:
                body = resp.json()
                msg = body.get("error") or body.get("message") or resp.text
            except Exception:
                msg = resp.text
            raise MagsError(msg, status_code=resp.status_code)
        if not resp.content:
            return {}
        return resp.json()

    # ── jobs ─────────────────────────────────────────────────────────

    def run(
        self,
        script: str,
        *,
        name: str | None = None,
        workspace_id: str | None = None,
        base_workspace_id: str | None = None,
        persistent: bool = False,
        no_sleep: bool = False,
        ephemeral: bool = False,
        startup_command: str | None = None,
        environment: Dict[str, str] | None = None,
        file_ids: List[str] | None = None,
        disk_gb: int | None = None,
        no_sync: bool = False,
    ) -> dict:
        """Submit a job for execution.

        Args:
            script: Shell script to execute inside the VM.
            name: Optional job name.
            workspace_id: Persistent workspace name (synced to S3).
            base_workspace_id: Read-only base workspace to mount.
            persistent: Keep VM alive after script finishes.
            no_sleep: Never auto-sleep this VM (requires persistent=True).
                The VM stays running 24/7 and auto-recovers if its host goes down.
            ephemeral: No S3 sync (faster, truly ephemeral).
            startup_command: Command to run when VM wakes from sleep.
            environment: Key-value env vars injected into the VM.
            file_ids: File IDs to download into VM before script runs.
            disk_gb: Custom disk size in GB (default 2GB).
            no_sync: Skip S3 sync, use local disk only.

        Returns ``{"request_id": ..., "status": "accepted"}``.
        """
        if ephemeral and workspace_id:
            raise MagsError("Cannot use ephemeral with workspace_id")
        if ephemeral and persistent:
            raise MagsError("Cannot use ephemeral with persistent")
        if no_sleep and not persistent:
            raise MagsError("no_sleep requires persistent=True")

        payload = {
            "script": script,
            "type": "inline",
            "persistent": persistent,
        }
        if no_sleep:
            payload["no_sleep"] = True
        if no_sync:
            payload["no_sync"] = True
        if name:
            payload["name"] = name
        if not ephemeral and workspace_id:
            payload["workspace_id"] = workspace_id
        if base_workspace_id:
            payload["base_workspace_id"] = base_workspace_id
        if startup_command:
            payload["startup_command"] = startup_command
        if environment:
            payload["environment"] = environment
        if file_ids:
            payload["file_ids"] = file_ids
        if disk_gb:
            payload["disk_gb"] = disk_gb

        return self._request("POST", "/mags-jobs", json=payload)

    def run_and_wait(
        self,
        script: str,
        *,
        timeout: float = 60.0,
        poll_interval: float = 1.0,
        **run_kwargs: Any,
    ) -> dict:
        """Submit a job and block until it completes or times out.

        Returns a dict with ``request_id``, ``status``, ``exit_code``,
        ``duration_ms``, and ``logs``.
        """
        result = self.run(script, **run_kwargs)
        request_id = result["request_id"]

        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            status = self.status(request_id)
            if status["status"] in ("completed", "error"):
                logs_resp = self.logs(request_id)
                return {
                    "request_id": request_id,
                    "status": status["status"],
                    "exit_code": status.get("exit_code", 0),
                    "duration_ms": status.get("script_duration_ms", 0),
                    "logs": logs_resp.get("logs", []),
                }
            time.sleep(poll_interval)

        raise MagsError(f"Job {request_id} timed out after {timeout}s")

    def status(self, request_id: str) -> dict:
        """Get the status of a job."""
        return self._request("GET", f"/mags-jobs/{request_id}/status")

    def logs(self, request_id: str) -> dict:
        """Get logs for a job. Returns ``{"logs": [...]}`."""
        return self._request("GET", f"/mags-jobs/{request_id}/logs")

    def list_jobs(self, *, page: int = 1, page_size: int = 20) -> dict:
        """List recent jobs. Returns ``{"jobs": [...], "total": N, ...}``."""
        return self._request(
            "GET", "/mags-jobs", params={"page": page, "page_size": page_size}
        )

    def update_job(
        self,
        request_id: str,
        *,
        startup_command: str | None = None,
        no_sleep: bool | None = None,
    ) -> dict:
        """Update a job's settings.

        Args:
            request_id: The job/workspace ID to update.
            startup_command: Command to run when VM wakes from sleep.
            no_sleep: If True, VM never auto-sleeps. If False, re-enables auto-sleep.
        """
        payload: dict = {}
        if startup_command is not None:
            payload["startup_command"] = startup_command
        if no_sleep is not None:
            payload["no_sleep"] = no_sleep
        return self._request("PATCH", f"/mags-jobs/{request_id}", json=payload)

    def enable_access(self, request_id: str, *, port: int = 8080) -> dict:
        """Enable external access (URL or SSH) for a persistent job's VM.

        Use ``port=22`` for SSH access, or ``port=8080`` (default) for HTTP/URL access.
        """
        return self._request(
            "POST", f"/mags-jobs/{request_id}/access", json={"port": port}
        )

    def stop(self, name_or_id: str) -> dict:
        """Stop a running job.

        Accepts a job ID, job name, or workspace ID.
        """
        request_id = self._resolve_job_id(name_or_id)
        return self._request("POST", f"/mags-jobs/{request_id}/stop")

    def resize(
        self,
        workspace: str,
        disk_gb: int,
        *,
        timeout: float = 30.0,
        poll_interval: float = 1.0,
    ) -> dict:
        """Resize a workspace's disk. Stops the existing VM, then creates a new one.

        Workspace files are preserved in S3.
        Returns ``{"request_id": ..., "status": "running"}``.
        """
        existing = self.find_job(workspace)
        if existing and existing.get("status") == "running":
            self._request("POST", f"/mags-jobs/{existing['request_id']}/sync")
            self._request("POST", f"/mags-jobs/{existing['request_id']}/stop")
            time.sleep(1)
        elif existing and existing.get("status") == "sleeping":
            self._request("POST", f"/mags-jobs/{existing['request_id']}/stop")
            time.sleep(1)

        return self.new(workspace, persistent=True, disk_gb=disk_gb,
                        timeout=timeout, poll_interval=poll_interval)

    def new(
        self,
        name: str,
        *,
        persistent: bool = False,
        base_workspace_id: str | None = None,
        disk_gb: int | None = None,
        timeout: float = 30.0,
        poll_interval: float = 1.0,
    ) -> dict:
        """Create a new VM sandbox and wait until it's running.

        By default, data lives on local disk only (no S3 sync).
        Pass ``persistent=True`` to enable S3 data persistence.

        Equivalent to ``mags new <name>`` or ``mags new <name> -p``.

        Returns ``{"request_id": ..., "status": "running"}``.
        """
        result = self.run(
            "sleep infinity",
            workspace_id=name,
            persistent=True,
            no_sync=not persistent,
            base_workspace_id=base_workspace_id,
            disk_gb=disk_gb,
        )
        request_id = result["request_id"]

        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            st = self.status(request_id)
            if st["status"] == "running" and st.get("vm_id"):
                return {"request_id": request_id, "status": "running"}
            if st["status"] in ("completed", "error"):
                raise MagsError(f"Job {request_id} ended unexpectedly: {st['status']}")
            time.sleep(poll_interval)

        raise MagsError(f"Job {request_id} did not start within {timeout}s")

    def find_job(self, name_or_id: str) -> dict | None:
        """Find a running or sleeping job by name, workspace ID, or job ID.

        Uses the same resolution priority as the CLI:
        running/sleeping exact name → workspace ID → any status exact name.
        Returns the job dict or ``None``.
        """
        jobs = self.list_jobs(page_size=50).get("jobs", [])

        # Priority 1: exact name match, running/sleeping
        for j in jobs:
            if j.get("name") == name_or_id and j.get("status") in ("running", "sleeping"):
                return j

        # Priority 2: workspace_id match, running/sleeping
        for j in jobs:
            if j.get("workspace_id") == name_or_id and j.get("status") in ("running", "sleeping"):
                return j

        # Priority 3: exact name match, any status
        for j in jobs:
            if j.get("name") == name_or_id:
                return j

        # Priority 4: workspace_id match, any status
        for j in jobs:
            if j.get("workspace_id") == name_or_id:
                return j

        return None

    def url(self, name_or_id: str, *, port: int = 8080) -> dict:
        """Enable public URL access for a job's VM.

        Accepts a job ID, job name, or workspace ID.
        Returns dict with ``url`` and access details.
        """
        request_id = self._resolve_job_id(name_or_id)
        st = self.status(request_id)
        resp = self.enable_access(request_id, port=port)
        subdomain = st.get("subdomain") or resp.get("subdomain")
        if subdomain:
            resp["url"] = f"https://{subdomain}.apps.magpiecloud.com"
        return resp

    def exec(self, name_or_id: str, command: str, *, timeout: int = 30) -> dict:
        """Execute a command on an existing running/sleeping sandbox via SSH.

        Equivalent to ``mags exec <workspace> '<command>'``.

        Returns ``{"exit_code": int, "output": str}``.
        """
        job = self.find_job(name_or_id)
        if not job:
            raise MagsError(f"No running or sleeping VM found for '{name_or_id}'")
        if job["status"] not in ("running", "sleeping"):
            raise MagsError(
                f"VM for '{name_or_id}' is {job['status']}, needs to be running or sleeping"
            )

        request_id = job.get("request_id") or job.get("id")

        # Call /access directly — for sleeping VMs this triggers the wake
        access = self.enable_access(request_id, port=22)

        if not access.get("success") or not access.get("ssh_host"):
            raise MagsError(
                f"Failed to enable SSH access: {access.get('error', 'unknown error')}"
            )

        ssh_host = access["ssh_host"]
        ssh_port = str(access["ssh_port"])
        ssh_key = access.get("ssh_private_key", "")

        # Wrap command to handle chroot overlay, same as CLI
        escaped = command.replace("'", "'\\''")
        wrapped = (
            f"if [ -d /overlay/bin ]; then "
            f"chroot /overlay /bin/sh -l -c 'cd /root 2>/dev/null; {escaped}'; "
            f"else cd /root 2>/dev/null; {escaped}; fi"
        )

        key_file = None
        try:
            ssh_args = [
                "ssh",
                "-o", "StrictHostKeyChecking=no",
                "-o", "UserKnownHostsFile=/dev/null",
                "-o", "LogLevel=ERROR",
                "-p", ssh_port,
            ]
            if ssh_key:
                fd, key_file = tempfile.mkstemp(prefix="mags_ssh_")
                os.write(fd, ssh_key.encode())
                os.close(fd)
                os.chmod(key_file, 0o600)
                ssh_args.extend(["-i", key_file])

            ssh_args.append(f"root@{ssh_host}")
            ssh_args.append(wrapped)

            proc = subprocess.run(
                ssh_args,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
            return {
                "exit_code": proc.returncode,
                "output": proc.stdout,
                "stderr": proc.stderr,
            }
        except subprocess.TimeoutExpired:
            raise MagsError(f"Command timed out after {timeout}s")
        finally:
            if key_file:
                try:
                    os.unlink(key_file)
                except OSError:
                    pass

    def usage(self, *, window_days: int = 30) -> dict:
        """Get aggregated usage summary."""
        return self._request(
            "GET", "/mags-jobs/usage", params={"window_days": window_days}
        )

    # ── internal helpers ─────────────────────────────────────────────

    def _resolve_job_id(self, name_or_id: str) -> str:
        """Resolve a job name, workspace ID, or UUID to a request_id."""
        # If it looks like a UUID, use directly
        if len(name_or_id) >= 32 and "-" in name_or_id:
            return name_or_id
        job = self.find_job(name_or_id)
        if not job:
            raise MagsError(f"No job found for '{name_or_id}'")
        return job.get("request_id") or job.get("id")

    # ── file uploads ─────────────────────────────────────────────────

    def upload_file(self, file_path: str) -> str:
        """Upload a single file. Returns the file ID."""
        p = Path(file_path)
        if not p.exists():
            raise MagsError(f"File not found: {file_path}")

        url = f"{self.api_url}/api/v1/mags-files"
        # Use a fresh request without the JSON content-type
        resp = requests.post(
            url,
            files={"file": (p.name, p.read_bytes(), "application/octet-stream")},
            headers={"Authorization": f"Bearer {self.api_token}"},
            timeout=self.timeout,
        )
        if resp.status_code >= 400:
            raise MagsError(resp.text, status_code=resp.status_code)
        data = resp.json()
        file_id = data.get("file_id")
        if not file_id:
            raise MagsError(f"Upload failed for {p.name}: {data}")
        return file_id

    def upload_files(self, file_paths: List[str]) -> List[str]:
        """Upload multiple files. Returns a list of file IDs."""
        return [self.upload_file(fp) for fp in file_paths]

    # ── workspaces ────────────────────────────────────────────────────

    def list_workspaces(self) -> dict:
        """List all workspaces. Returns ``{"workspaces": [...], "total": N}``."""
        return self._request("GET", "/mags-workspaces")

    def delete_workspace(self, workspace_id: str) -> dict:
        """Delete a workspace and all its stored data.

        This permanently removes the workspace filesystem from S3.
        Active jobs using the workspace must be stopped first.
        """
        return self._request("DELETE", f"/mags-workspaces/{workspace_id}")

    def sync(self, request_id: str) -> dict:
        """Sync a running job's workspace to S3 without stopping the VM.

        Use this to persist workspace changes immediately, e.g. after
        setting up a base image.
        """
        return self._request("POST", f"/mags-jobs/{request_id}/sync")

    # ── cron jobs ────────────────────────────────────────────────────

    def cron_create(
        self,
        *,
        name: str,
        cron_expression: str,
        script: str,
        workspace_id: str | None = None,
        environment: Dict[str, str] | None = None,
        persistent: bool = False,
    ) -> dict:
        """Create a scheduled cron job."""
        payload = {
            "name": name,
            "cron_expression": cron_expression,
            "script": script,
            "persistent": persistent,
        }
        if workspace_id:
            payload["workspace_id"] = workspace_id
        if environment:
            payload["environment"] = environment
        return self._request("POST", "/mags-cron", json=payload)

    def cron_list(self) -> dict:
        """List all cron jobs."""
        return self._request("GET", "/mags-cron")

    def cron_get(self, cron_id: str) -> dict:
        """Get a cron job by ID."""
        return self._request("GET", f"/mags-cron/{cron_id}")

    def cron_update(self, cron_id: str, **updates: Any) -> dict:
        """Update a cron job. Pass fields as keyword arguments."""
        return self._request("PATCH", f"/mags-cron/{cron_id}", json=updates)

    def cron_delete(self, cron_id: str) -> dict:
        """Delete a cron job."""
        return self._request("DELETE", f"/mags-cron/{cron_id}")

    # ── URL aliases ──────────────────────────────────────────────────

    def url_alias_create(
        self,
        subdomain: str,
        workspace_id: str,
        domain: str = "apps.magpiecloud.com",
    ) -> dict:
        """Create a stable URL alias for a workspace.

        The alias maps ``subdomain.<domain>`` to the active job in the workspace.
        Use ``domain="app.lfg.run"`` for the LFG domain.

        Returns ``{"id": ..., "subdomain": ..., "url": ...}``.
        """
        return self._request(
            "POST",
            "/mags-url-aliases",
            json={
                "subdomain": subdomain,
                "workspace_id": workspace_id,
                "domain": domain,
            },
        )

    def url_alias_list(self) -> dict:
        """List all URL aliases. Returns ``{"aliases": [...], "total": N}``."""
        return self._request("GET", "/mags-url-aliases")

    def url_alias_delete(self, subdomain: str) -> dict:
        """Delete a URL alias by subdomain."""
        return self._request("DELETE", f"/mags-url-aliases/{subdomain}")
