"""IsardVDI API client for managing virtual desktops.

Mirrors the PowerShell functions in desktop.ps1:
- IsardVdiClient.login()  → Get-IsardToken
- get_desktops()          → Get-Isard

Authentication uses the IsardVDI SAML provider, which triggers the same
Azure Entra ID flow as the Gencat SAML login.  The SAML assertion is posted
to IsardVDI's own ACS endpoint, which returns (or sets) a JWT token.
"""

from __future__ import annotations

import logging
import re
import time
from typing import Optional
from urllib.parse import urlparse, parse_qs

import httpx

from .auth import (
    AuthFlow,
    _extract_form_field,
    _extract_form_action,
    _extract_hidden_form,
    _parse_azure_config,
)
from .client import GencatClient
from .errors import (
    AzureAuthError,
    AzureRedirectError,
    BookingRequiredError,
    HttpError,
    NetworkError,
)
from .types import BastionTarget, Desktop, GencatConfig, IsardVdiSession

logger = logging.getLogger(__name__)

_BASE_URL = "https://elmeuescriptori.gestioeducativa.gencat.cat"

# Hardcoded in desktop.ps1 — the "Provençana" category on this IsardVDI instance.
_DEFAULT_CATEGORY_ID = "db79b78a-5408-4ff9-9853-a1172b7eadb2"


def _extract_jwt_from_response(
    url: str,
    body: str,
    cookies: dict,
) -> Optional[str]:
    """Try to extract a JWT token from a completed IsardVDI SAML ACS response.

    IsardVDI can return the token in several ways:
    1. Raw JWT string as the response body  ← checked first (most reliable)
    2. JSON body: {"token": "...", ...}
    3. Cookie named ``token``, ``isardvdi_session``, or ``jwt``
    4. Redirect URL with ``?token=...`` query parameter
       (``state`` is intentionally excluded — it holds a transient callback JWT)
    """
    # 1. Check response body — raw JWT (three dot-separated base64url segments)
    stripped = body.strip().strip('"')
    if re.match(r"^[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+$", stripped):
        logger.debug("JWT found as raw response body")
        return stripped

    # 2. Check body as JSON {"token": "..."}
    try:
        import json as _json

        data = _json.loads(body)
        if isinstance(data, dict):
            for key in ("token", "jwt", "access_token"):
                if key in data and data[key]:
                    tok = str(data[key]).strip().strip('"')
                    if tok:
                        logger.debug("JWT found in JSON body key %r", key)
                        return tok
    except Exception:
        pass

    # 3. Check cookies
    for key in ("token", "isardvdi_session", "isardvdi", "jwt", "access_token"):
        if key in cookies:
            tok = str(cookies[key]).strip().strip('"')
            if tok:
                logger.debug("JWT found in cookie %r", key)
                return tok

    # 4. Check URL query params (excluding 'state' which is a transient callback token)
    parsed = urlparse(url)
    qs = parse_qs(parsed.query)
    for key in ("token", "jwt", "access_token"):
        if key in qs:
            tok = qs[key][0].strip().strip('"')
            if tok:
                logger.debug("JWT found in URL query param %r", key)
                return tok

    return None


class _IsardVdiAuthFlow(AuthFlow):
    """SAML auth flow variant that targets IsardVDI's SAML endpoints.

    Overrides the ACS submission step to post the SAMLResponse to IsardVDI's
    own ACS URL and extract the resulting JWT token instead of building a
    Gencat :class:`AuthSession`.
    """

    def __init__(
        self,
        client: GencatClient,
        config: GencatConfig,
        category_id: str,
    ):
        super().__init__(client, config)
        self._category_id = category_id
        self._isardvdi_jwt: Optional[str] = None

    def login_isardvdi(self, username: str, password: str) -> str:
        """Run the full SAML flow and return the IsardVDI JWT token.

        Raises:
            :class:`AzureAuthError`: If credentials are rejected or the SAML
                flow fails to produce a token.
        """
        base = self._config.base_url.rstrip("/")
        category_id = self._category_id

        # ── Step 1: Initiate IsardVDI SAML login ─────────────────────────
        logger.info("IsardVDI SAML Step 1: initiating login")
        login_url = (
            f"{base}/authentication/saml/login?provider=saml&category_id={category_id}"
        )
        azure_url, _ = self._client.get_following_redirects(login_url)

        if "login.microsoftonline.com" not in azure_url:
            raise AzureRedirectError(
                f"Expected redirect to Azure login, got: {azure_url}",
                redirect_url=azure_url,
            )

        # Capture the original RelayState from the Azure redirect URL.
        # This opaque token matches the saml_* cookie set by IsardVDI and
        # MUST be preserved through the Azure hop loop — the final form page
        # may contain a corrupted/re-encoded RelayState value.
        _azure_qs = parse_qs(urlparse(azure_url).query)
        original_relay_state: str = _azure_qs.get("RelayState", [""])[0]
        logger.debug("Captured original RelayState: %r", original_relay_state)

        # ── Steps 2-4: Standard Azure credential submission ───────────────
        # Reuse the parent class helpers directly.
        azure_landing_url, azure_html = self._client.get_following_redirects(azure_url)

        # BssoInterrupt handling (same as parent)
        for _ in range(3):
            if "BssoInterrupt" not in azure_html:
                break
            logger.debug("BssoInterrupt detected, re-fetching Azure login page")
            new_url, new_html = self._client.get_following_redirects(azure_url)
            if "BssoInterrupt" in new_html:
                try:
                    cfg_tmp = _parse_azure_config(azure_html)
                    bsso_url = cfg_tmp.url_post.replace(r"\u0026", "&").replace(
                        r"\u003d", "="
                    )
                    new_url, new_html = self._client.get_following_redirects(bsso_url)
                except Exception:
                    pass
            azure_landing_url = new_url
            azure_html = new_html

        azure_config = _parse_azure_config(azure_html)

        flow_token = self._get_credential_type(username, azure_config)
        post_url, post_body = self._submit_credentials(
            username, password, flow_token, azure_config
        )

        # ── Post-credential redirect loop ─────────────────────────────────
        for hop in range(15):
            if "SAMLResponse" in post_body:
                break
            if "oPostParams" in post_body:
                logger.info("Hop %d: following oPostParams", hop)
                post_url, post_body = self._follow_o_post_params(
                    post_body, azure_config.correlation_id
                )
                continue
            method = ""
            if "<form" in post_body.lower():
                m = re.search(
                    r'<form[^>]+method=["\'](\w+)["\']', post_body, re.IGNORECASE
                )
                method = m.group(1).upper() if m else "GET"
            if "<form" in post_body.lower() and method == "POST":
                result = _extract_hidden_form(post_body, post_url)
                if result is not None:
                    action, fields = result
                    if base in action:
                        break
                    logger.info("Hop %d: following hidden form to %s", hop, action)
                    post_url, post_body = self._client.post_form_following_redirects(
                        action, fields
                    )
                    continue
            if "$Config=" in post_body:
                try:
                    cfg = _parse_azure_config(post_body)
                except Exception:
                    break
                if cfg.url_post:
                    logger.info(
                        "Hop %d: KmsiInterrupt, posting to %s", hop, cfg.url_post
                    )
                    kmsi_form = {
                        "LoginOptions": "3",
                        "type": "28",
                        "ctx": cfg.s_ctx,
                        "flowToken": cfg.flow_token,
                        "canary": cfg.canary,
                        "hpgrequestid": cfg.session_id,
                    }
                    post_url, post_body = self._client.post_form_following_redirects(
                        cfg.url_post, kmsi_form
                    )
                    continue
            break

        # ── Submit SAMLResponse to IsardVDI ACS ───────────────────────────
        if "SAMLResponse" not in post_body:
            raise AzureAuthError(
                f"SAMLResponse not found after Azure auth loop; last URL: {post_url}",
                correlation_id=azure_config.correlation_id,
            )

        saml_response = _extract_form_field(post_body, "SAMLResponse")
        if not saml_response:
            raise AzureAuthError("Could not extract SAMLResponse field from page")

        form_action = (
            _extract_form_action(post_body) or f"{base}/authentication/saml/acs"
        )

        # Use the original RelayState captured at step 1, NOT the one embedded
        # in the final Azure form — by this point the hop chain (KMSI, OAuth2Cl,
        # Terms-of-Use) may have re-encoded the SAMLResponse inside RelayState.
        relay_state_for_acs = original_relay_state or (
            _extract_form_field(post_body, "RelayState") or ""
        )
        logger.debug("Using RelayState for ACS: %r", relay_state_for_acs[:60])

        logger.info("IsardVDI SAML Step 5: posting SAMLResponse to %s", form_action)
        form = {"SAMLResponse": saml_response, "RelayState": relay_state_for_acs}

        # We need access to the underlying httpx client's cookie jar for this request.
        acs_url, acs_body = self._client.post_form_following_redirects(
            form_action, form
        )
        cookies = dict(self._client._client.cookies)
        logger.debug("ACS response URL: %s", acs_url)
        logger.debug("ACS response body prefix: %r", acs_body[:80])
        logger.debug("ACS cookie keys: %r", list(cookies.keys()))

        # Extract JWT from response
        token = _extract_jwt_from_response(acs_url, acs_body, cookies)
        if not token:
            raise AzureAuthError(
                f"IsardVDI SAML login did not return a JWT token. ACS URL={acs_url!r}"
            )

        return token


class IsardVdiClient:
    """HTTP client for the IsardVDI REST API (``/api/v3``)."""

    def __init__(self, session: IsardVdiSession, base_url: str = _BASE_URL):
        self._base_url = base_url.rstrip("/")
        self._session = session
        self._client = httpx.Client(timeout=30.0, verify=True)

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "IsardVdiClient":
        return self

    def __exit__(self, *_) -> None:
        self.close()

    # ------------------------------------------------------------------
    # Authentication  (mirrors Get-IsardToken in desktop.ps1)
    # ------------------------------------------------------------------

    @classmethod
    def login(
        cls,
        username: str,
        password: str,
        base_url: str = _BASE_URL,
        category_id: str = _DEFAULT_CATEGORY_ID,
    ) -> "IsardVdiClient":
        """Authenticate against IsardVDI via the SAML provider (Azure Entra ID).

        Triggers the ``saml`` provider login flow starting at
        ``/authentication/login?provider=saml&category_id=<id>``, which
        redirects to Azure, completes the credential challenge, and posts
        the resulting SAMLResponse to IsardVDI's ACS.  The JWT token
        returned by the ACS is used for subsequent API calls.

        Args:
            username: Gencat username (e.g. ``user@edu.gencat.cat``).
            password: Gencat password.
            base_url: IsardVDI instance base URL.
            category_id: IsardVDI category UUID (defaults to the Provençana category).

        Returns:
            An :class:`IsardVdiClient` ready to call API endpoints.

        Raises:
            :class:`AzureAuthError`: Login rejected or SAML flow failed.
            :class:`HttpError`: Other non-2xx response.
            :class:`NetworkError`: Network-level connectivity error.
        """
        config = GencatConfig(base_url=base_url)
        with GencatClient(config) as http_client:
            flow = _IsardVdiAuthFlow(http_client, config, category_id)
            token = flow.login_isardvdi(username, password)

        session = IsardVdiSession(token=token, username=username)
        return cls(session, base_url=base_url)

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _headers(self) -> dict[str, str]:
        return {
            "Accept": "application/json",
            "Authorization": f"Bearer {self._session.token}",
        }

    def _get(self, endpoint: str) -> dict | list:
        url = f"{self._base_url}{endpoint}"
        try:
            resp = self._client.get(url, headers=self._headers())
        except httpx.RequestError as exc:
            raise NetworkError(f"Request failed: {exc}", exc) from exc

        if resp.status_code >= 400:
            raise HttpError(resp.status_code, f"GET {endpoint}")

        return resp.json()

    def _put(self, endpoint: str, body: dict) -> dict | list:
        url = f"{self._base_url}{endpoint}"
        headers = {**self._headers(), "Content-Type": "application/json"}
        try:
            resp = self._client.put(url, headers=headers, json=body)
        except httpx.RequestError as exc:
            raise NetworkError(f"Request failed: {exc}", exc) from exc

        if resp.status_code >= 400:
            raise HttpError(resp.status_code, f"PUT {endpoint}")

        return resp.json()

    # ------------------------------------------------------------------
    # Desktop operations  (mirror of desktop.ps1 Get-Isard / Start-Isard /
    #                       Stop-Isard)
    # ------------------------------------------------------------------

    def get_desktops(self, name: Optional[str] = None) -> list[Desktop]:
        """Return all desktops for the authenticated user.

        Args:
            name: If provided, filter to the single desktop with this name.

        Returns:
            A list of :class:`Desktop` objects sorted by name.

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        data = self._get("/api/v3/user/desktops")
        if not isinstance(data, list):
            data = [data]

        desktops = [Desktop.from_dict(item) for item in data]
        desktops.sort(key=lambda d: d.name)

        if name is not None:
            desktops = [d for d in desktops if d.name == name]

        return desktops

    def stop_desktop(self, desktop_id: str) -> str:
        """Stop a running desktop.

        Mirrors ``Stop-Isard`` in desktop.ps1.

        Args:
            desktop_id: The desktop UUID (``Desktop.id``).

        Returns:
            The new desktop state string as reported by the API (e.g.
            ``"Stopping"``, ``"Shutting-down"``, or ``"Stopped"``).

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        data = self._get(f"/api/v3/desktop/stop/{desktop_id}")
        return data.get("status", "") if isinstance(data, dict) else ""

    def start_desktop(self, desktop_id: str) -> str:
        """Start a stopped desktop.

        Mirrors ``Start-Isard`` in desktop.ps1.

        Args:
            desktop_id: The desktop UUID (``Desktop.id``).

        Returns:
            The new desktop state string as reported by the API (e.g.
            ``"Starting"`` or ``"Started"``).

        Raises:
            :class:`BookingRequiredError`: HTTP 428 — a booking must be made before
                the desktop can be started.
            :class:`HttpError`: Other non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        try:
            data = self._get(f"/api/v3/desktop/start/{desktop_id}")
        except HttpError as exc:
            if exc.status_code == 428:
                raise BookingRequiredError(desktop_id) from exc
            raise
        return data.get("status", "") if isinstance(data, dict) else ""

    def star_desktop(self, desktop_id: str, starred: bool = True) -> None:
        """Mark (or unmark) a desktop as a favourite / starred.

        Uses ``PUT /api/v3/domain/{id}`` with ``{"favourite": <bool>}``.

        Args:
            desktop_id: The desktop UUID (``Desktop.id``).
            starred: ``True`` to star the desktop, ``False`` to unstar it.

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        self._put(f"/api/v3/domain/{desktop_id}", {"favourite": starred})

    def wait_for_state(
        self,
        desktop_id: str,
        target_state: str,
        timeout: float = 120.0,
        poll_interval: float = 3.0,
    ) -> str:
        """Poll the desktop state until it matches *target_state* or the timeout expires.

        Args:
            desktop_id: The desktop UUID to poll.
            target_state: The state to wait for, e.g. ``"Started"`` or ``"Stopped"``.
            timeout: Maximum seconds to wait before raising :class:`TimeoutError`.
            poll_interval: Seconds between polls.

        Returns:
            The final desktop state string.

        Raises:
            :class:`TimeoutError`: Desktop did not reach *target_state* within *timeout* seconds.
            :class:`HttpError`: Non-2xx response from the API during polling.
            :class:`NetworkError`: Network-level connectivity error.
        """
        deadline = time.monotonic() + timeout
        while True:
            desktops = self.get_desktops()
            match = next((d for d in desktops if d.id == desktop_id), None)
            current_state = match.state if match else "Unknown"
            if current_state == target_state:
                return current_state
            if time.monotonic() >= deadline:
                raise TimeoutError(
                    f"Desktop {desktop_id!r} did not reach state {target_state!r} "
                    f"within {timeout:.0f}s (last state: {current_state!r})"
                )
            time.sleep(poll_interval)

    @property
    def isardvdi_session(self) -> IsardVdiSession:
        return self._session

    # ------------------------------------------------------------------
    # Bastion (SSH jump-host)  — /api/v3/desktop/bastion and related
    # ------------------------------------------------------------------

    def is_bastion_allowed(self) -> bool:
        """Return ``True`` if the current user is allowed to use the SSH bastion.

        Calls ``GET /api/v3/user/bastion/allowed``.

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        data = self._get("/api/v3/user/bastion/allowed")
        if isinstance(data, bool):
            return data
        if isinstance(data, dict):
            return bool(data.get("allowed", data.get("result", False)))
        return bool(data)

    def get_bastion(self, desktop_id: str) -> BastionTarget:
        """Return the current bastion configuration for a desktop.

        Calls ``GET /api/v3/desktop/bastion/{desktop_id}``.

        Args:
            desktop_id: The desktop UUID.

        Returns:
            A :class:`BastionTarget` with the current SSH/HTTP bastion settings.

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        data = self._get(f"/api/v3/desktop/bastion/{desktop_id}")
        if not isinstance(data, dict):
            raise HttpError(
                500, f"Unexpected response type from bastion endpoint: {type(data)}"
            )
        return BastionTarget.from_dict(data)

    def set_bastion(
        self,
        desktop_id: str,
        *,
        ssh_enabled: bool = True,
        ssh_port: int = 22,
        authorized_keys: Optional[list[str]] = None,
        http_enabled: bool = False,
    ) -> BastionTarget:
        """Enable (or update) the SSH bastion for a desktop.

        Fetches the current config, merges the requested changes, and PUTs the
        result back.  Only the fields you pass are changed; everything else is
        left as-is.

        Args:
            desktop_id: The desktop UUID.
            ssh_enabled: Whether SSH bastion access should be active.
            ssh_port: The VM-side SSH port (default 22 — the port inside the VM,
                **not** the bastion listen port).
            authorized_keys: If provided, *replaces* the full key list for this
                desktop.  Pass an empty list to clear all keys.  ``None`` leaves
                the existing keys unchanged.
            http_enabled: Whether HTTP proxying should be active.

        Returns:
            The updated :class:`BastionTarget` (re-fetched after the PUT).

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        current = self.get_bastion(desktop_id)

        keys = (
            authorized_keys
            if authorized_keys is not None
            else current.ssh.authorized_keys
        )

        body = {
            "ssh": {
                "enabled": ssh_enabled,
                "port": ssh_port,
                "authorized_keys": keys,
            },
            "http": {
                "enabled": http_enabled,
                "http_port": current.http.http_port,
                "https_port": current.http.https_port,
                "proxy_protocol": current.http.proxy_protocol,
            },
            "domains": current.domains,
        }
        self._put(f"/api/v3/desktop/bastion/{desktop_id}", body)
        return self.get_bastion(desktop_id)

    def add_ssh_key(self, desktop_id: str, public_key: str) -> BastionTarget:
        """Add a single SSH public key to a desktop's authorized_keys list.

        Idempotent — if the key is already present it is not duplicated.

        Args:
            desktop_id: The desktop UUID.
            public_key: The full public-key string, e.g. ``"ssh-ed25519 AAAA... user@host"``.

        Returns:
            The updated :class:`BastionTarget`.

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        current = self.get_bastion(desktop_id)
        existing = current.ssh.authorized_keys
        key = public_key.strip()
        if key not in existing:
            existing = existing + [key]
        return self.set_bastion(
            desktop_id,
            ssh_enabled=current.ssh.enabled,
            ssh_port=current.ssh.port,
            authorized_keys=existing,
            http_enabled=current.http.enabled,
        )

    @property
    def bastion_host(self) -> str:
        """The SSH hostname for the bastion (same as the IsardVDI instance host)."""
        from urllib.parse import urlparse as _up

        return _up(self._base_url).hostname or self._base_url

    @property
    def bastion_port(self) -> int:
        """The SSH listen port for the bastion.

        IsardVDI defaults to ``BASTION_SSH_PORT`` env var, which itself falls
        back to ``HTTPS_PORT`` (443).  Port 22 is only used on custom
        deployments that override both variables.
        """
        return 443

    def ssh_command(self, target: BastionTarget, ssh_user: str = "user") -> list[str]:
        """Return the ``ssh`` argv list to connect to *target* via the bastion.

        The bastion uses the target's ``id`` as the SSH username to the jump
        host, which then proxies the connection into the VM.

        Connection format::

            ssh -J <target.id>@<bastion_host>:<bastion_port> <ssh_user>@<bastion_host>

        Because IsardVDI's bastion is a true SSH proxy (not just a jump host),
        the actual command is simpler::

            ssh <target.id>@<bastion_host> -p <bastion_port>

        Args:
            target: The :class:`BastionTarget` for the desktop.
            ssh_user: The Linux username inside the VM (default ``"user"``).

        Returns:
            A list of strings suitable for :func:`subprocess.exec` / :func:`os.execvp`.
        """
        cmd = [
            "ssh",
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            f"{target.id}@{self.bastion_host}",
        ]
        if self.bastion_port != 22:
            cmd += ["-p", str(self.bastion_port)]
        return cmd

    # ------------------------------------------------------------------
    # Template
    # ------------------------------------------------------------------

    def get_template_hardware(self, template_id: str) -> dict:
        """Return the hardware dict for a template.

        Calls ``GET /api/v3/template/{template_id}`` and extracts
        ``create_dict.hardware``.  Returns an empty dict on any error so
        callers can degrade gracefully.

        Args:
            template_id: The template UUID.

        Returns:
            A dict with at least ``vcpus`` (int) and ``memory`` (int, KiB)
            when available, otherwise ``{}``.

        Raises:
            :class:`HttpError`: Non-2xx response from the API.
            :class:`NetworkError`: Network-level connectivity error.
        """
        data = self._get(f"/api/v3/template/{template_id}")
        if not isinstance(data, dict):
            return {}
        return data.get("create_dict", {}).get("hardware", {})

    def enrich_hardware(self, desktops: list) -> list:
        """Populate ``vcpus`` and ``memory_mb`` on each desktop by fetching templates.

        Fetches each unique template once and fills in the hardware fields
        in-place.  Desktops without a template or whose template fetch fails
        are left with ``None`` values.

        Args:
            desktops: List of :class:`Desktop` objects (mutated in-place).

        Returns:
            The same list, with ``vcpus`` and ``memory_mb`` filled where
            available.
        """
        # Collect unique template IDs
        template_ids = {d.template_id for d in desktops if d.template_id}
        hw_cache: dict[str, dict] = {}
        for tid in template_ids:
            try:
                hw_cache[tid] = self.get_template_hardware(tid)
            except Exception:
                hw_cache[tid] = {}

        for d in desktops:
            hw = hw_cache.get(d.template_id or "", {})
            if "vcpus" in hw:
                d.vcpus = int(hw["vcpus"])
            if "memory" in hw:
                # API returns KiB; convert to MiB
                d.memory_mb = int(hw["memory"]) // 1024

        return desktops
