"""Full SAML 2.0 + Azure Entra ID authentication flow.

Mirrors auth/mod.rs from the Rust isard crate.

Flow overview
─────────────
1. GET  {base_url}/authentication/saml/login?provider=saml
       → follows redirects → lands on login.microsoftonline.com
2. GET  Azure login URL → extract $Config JSON
       (handles BssoInterrupt loops up to 3 hops)
3. POST GetCredentialType → verify user, get updated flowToken
4. POST credentials → multi-hop intermediate page loop (up to 15 hops):
       a. oPostParams (second BssoInterrupt)
       b. Hidden auto-submit <form>
       c. $Config-based page (KmsiInterrupt / "Keep me signed in")
5. POST SAMLResponse to Gencat ACS
6. Parse & validate SAML assertion → return AuthSession
"""

from __future__ import annotations

import json
import logging
import re
from datetime import datetime, timedelta, timezone
from typing import Optional
from urllib.parse import urljoin

from bs4 import BeautifulSoup

from .client import GencatClient
from .errors import (
    AzureAuthError,
    AzureFormParsingError,
    AzureRedirectError,
    SamlAssertionError,
)
from .saml import parse_saml_response_base64, validate_response
from .types import AuthSession, GencatConfig

logger = logging.getLogger(__name__)

_AZURE_BASE = "https://login.microsoftonline.com"
_GCT_URL = "https://login.microsoftonline.com/common/GetCredentialType?mkt=en-US"


# ---------------------------------------------------------------------------
# Internal data class for parsed $Config
# ---------------------------------------------------------------------------


class _AzureConfig:
    __slots__ = (
        "url_post",
        "flow_token",
        "api_canary",
        "canary",
        "correlation_id",
        "s_ctx",
        "session_id",
    )

    def __init__(
        self,
        url_post: str,
        flow_token: str,
        api_canary: str,
        canary: str,
        correlation_id: str,
        s_ctx: str,
        session_id: str,
    ):
        self.url_post = url_post
        self.flow_token = flow_token
        self.api_canary = api_canary
        self.canary = canary
        self.correlation_id = correlation_id
        self.s_ctx = s_ctx
        self.session_id = session_id


# ---------------------------------------------------------------------------
# $Config extraction helper
# ---------------------------------------------------------------------------


def _extract_dollar_config(html: str) -> dict:
    """Extract and parse the $Config={...} JSON block from an Azure page."""
    marker = "$Config="
    start = html.find(marker)
    if start == -1:
        raise AzureFormParsingError(
            "Could not find $Config block in Azure login page",
            html_snippet=html[:500],
        )

    brace_start = html.find("{", start)
    if brace_start == -1:
        raise AzureFormParsingError("Could not find opening { in $Config")

    # Walk characters to find matching closing brace
    depth = 0
    in_string = False
    escape_next = False
    brace_end = brace_start

    for i, ch in enumerate(html[brace_start:], start=brace_start):
        if escape_next:
            escape_next = False
            continue
        if ch == "\\" and in_string:
            escape_next = True
            continue
        if ch == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                brace_end = i + 1
                break

    json_str = html[brace_start:brace_end]
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as exc:
        raise AzureFormParsingError(
            f"Failed to parse $Config JSON: {exc}",
            html_snippet=json_str[:200],
        ) from exc


def _parse_azure_config(html: str) -> _AzureConfig:
    v = _extract_dollar_config(html)

    url_post_raw = v.get("urlPost", "")
    if not url_post_raw:
        raise AzureFormParsingError("Missing urlPost in $Config")

    url_post = (
        url_post_raw
        if url_post_raw.startswith("http")
        else f"{_AZURE_BASE}{url_post_raw}"
    )

    return _AzureConfig(
        url_post=url_post,
        flow_token=v.get("sFT", ""),
        api_canary=v.get("apiCanary", ""),
        canary=v.get("canary", ""),
        correlation_id=v.get("correlationId", ""),
        s_ctx=v.get("sCtx", ""),
        session_id=v.get("sessionId", ""),
    )


# ---------------------------------------------------------------------------
# HTML helpers
# ---------------------------------------------------------------------------


def _extract_hidden_form(
    html: str, page_url: str
) -> Optional[tuple[str, dict[str, str]]]:
    """Find the first <form> and collect all hidden/untyped inputs.

    Returns (action_url, {name: value}) or None.
    """
    soup = BeautifulSoup(html, "lxml")
    form = soup.find("form")
    if form is None:
        return None

    action_raw = form.get("action", "")
    action = urljoin(page_url, action_raw) if action_raw else page_url

    fields: dict[str, str] = {}
    for inp in form.find_all("input"):
        typ = (inp.get("type") or "").lower()
        if typ in ("hidden", ""):
            name = inp.get("name", "")
            value = inp.get("value", "")
            if name:
                fields[name] = value

    return action, fields


def _extract_form_field(html: str, field_name: str) -> Optional[str]:
    """Extract a named input's value from HTML."""
    esc = re.escape(field_name)
    patterns = [
        rf'<input[^>]+name=["\'{esc}["\'][^>]*value=["\']([^"\']*)["\']',
        rf'<input[^>]+value=["\']([^"\']*)["\'][^>]*name=["\'{esc}["\']',
    ]
    for pat in patterns:
        m = re.search(pat, html, re.IGNORECASE)
        if m:
            return m.group(1)
    return None


def _extract_form_action(html: str) -> Optional[str]:
    """Extract the action URL from the first <form> tag."""
    m = re.search(r'<form[^>]+action=["\']([^"\']+)["\']', html, re.IGNORECASE)
    if not m:
        return None
    return (
        m.group(1)
        .replace("&amp;", "&")
        .replace("&lt;", "<")
        .replace("&gt;", ">")
        .replace("&quot;", '"')
    )


def _extract_azure_error(html: str, correlation_id: str = "") -> None:
    """Raise AzureAuthError if an AADSTS error or wrong-password indicator is found."""
    m = re.search(r"AADSTS(\d+)", html)
    if m:
        code = f"AADSTS{m.group(1)}"
        raise AzureAuthError(
            f"Azure authentication error {code}",
            error_code=code,
            correlation_id=correlation_id or None,
        )
    if any(s in html for s in ("sErrorCode", "sErrTxt", "idTd_Pwd_Error")):
        raise AzureAuthError(
            "Invalid username or password",
            correlation_id=correlation_id or None,
        )


# ---------------------------------------------------------------------------
# Authentication flow
# ---------------------------------------------------------------------------


class AuthFlow:
    def __init__(self, client: GencatClient, config: GencatConfig):
        self._client = client
        self._config = config

    def authenticate(self, username: str, password: str) -> AuthSession:
        # ── Step 1: Initiate SAML login ───────────────────────────────────
        logger.info("Step 1: Initiating SAML login")
        login_url = f"{self._config.base_url}/authentication/saml/login?provider=saml"
        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,
            )

        # ── Step 2: Fetch Azure login page, extract $Config ───────────────
        logger.info("Step 2: Fetching Azure login page")
        azure_landing_url, azure_html = self._client.get_following_redirects(azure_url)

        # Handle BssoInterrupt (desktop/Kerberos SSO unavailable)
        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:
                # Still on BssoInterrupt — try following urlPost
                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 AzureFormParsingError:
                    pass

            azure_landing_url = new_url
            azure_html = new_html

        azure_config = _parse_azure_config(azure_html)
        logger.debug(
            "Parsed Azure config, correlationId=%s", azure_config.correlation_id
        )

        # ── Step 3: Submit username to GetCredentialType ──────────────────
        logger.info("Step 3: Submitting username")
        flow_token = self._get_credential_type(username, azure_config)
        logger.debug("Received updated flowToken")

        # ── Step 4: Submit password ───────────────────────────────────────
        logger.info("Step 4: Submitting password")
        post_url, post_body = self._submit_credentials(
            username, password, flow_token, azure_config
        )

        # ── Step 4b: 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 (BssoInterrupt)", hop)
                post_url, post_body = self._follow_o_post_params(
                    post_body, azure_config.correlation_id
                )
                continue

            # Hidden auto-submit form
            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
                    # Stop if the form targets the Gencat ACS directly
                    if self._config.base_url 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

            # $Config-based page (KmsiInterrupt, etc.)
            if "$Config=" in post_body:
                try:
                    cfg = _parse_azure_config(post_body)
                except AzureFormParsingError:
                    break
                if cfg.url_post:
                    logger.info(
                        "Hop %d: KmsiInterrupt / $Config page, posting to %s",
                        hop,
                        cfg.url_post,
                    )
                    kmsi_form = {
                        "LoginOptions": "3",  # Don't keep me signed in
                        "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

            # Nothing to follow
            break

        # ── Step 5: Submit SAMLResponse to Gencat ACS ─────────────────────
        logger.info("Step 5: Processing SAML ACS response")
        if "SAMLResponse" in post_body or "/authentication/saml/acs" in post_url:
            return self._process_acs_response(post_url, post_body)

        if (
            self._config.base_url in post_url
            and "saml/login" not in post_url
            and "SAMLResponse" in post_body
        ):
            return self._process_acs_response(post_url, post_body)

        raise AzureAuthError(
            f"Unexpected post-login URL: {post_url}",
            correlation_id=azure_config.correlation_id,
        )

    # ------------------------------------------------------------------
    # Step 3: GetCredentialType
    # ------------------------------------------------------------------

    def _get_credential_type(self, username: str, config: _AzureConfig) -> str:
        payload = json.dumps(
            {
                "username": username,
                "isOtherIdpSupported": True,
                "checkPhones": False,
                "isRemoteNGCSupported": True,
                "isCookieBannerShown": False,
                "isFidoSupported": True,
                "originalRequest": "",
                "country": "ES",
                "forceotclogin": False,
                "isExternalFederationDisallowed": False,
                "isRemoteConnectSupported": False,
                "federationFlags": 0,
                "isSignup": False,
                "flowToken": config.flow_token,
                "isAccessPassSupported": True,
            }
        )
        response_text = self._client.post_json(_GCT_URL, payload)
        try:
            v = json.loads(response_text)
        except json.JSONDecodeError as exc:
            raise AzureFormParsingError(
                f"Failed to parse GetCredentialType response: {exc}",
                html_snippet=response_text[:200],
            ) from exc

        if_exists = v.get("IfExistsResult")
        if if_exists is not None and if_exists != 0:
            raise AzureAuthError(
                f"User not found (IfExistsResult={if_exists})",
                correlation_id=config.correlation_id,
            )

        new_token = v.get("FlowToken", "") or ""
        return new_token if new_token else config.flow_token

    # ------------------------------------------------------------------
    # Step 4: Submit credentials
    # ------------------------------------------------------------------

    def _submit_credentials(
        self,
        username: str,
        password: str,
        flow_token: str,
        config: _AzureConfig,
    ) -> tuple[str, str]:
        form = {
            "login": username,
            "loginfmt": username,
            "passwd": password,
            "flowToken": flow_token,
            "canary": config.canary,
            "ctx": config.s_ctx,
            "hpgrequestid": config.session_id,
            "type": "11",
            "PPFT": flow_token,
        }
        url, body = self._client.post_form_following_redirects(config.url_post, form)
        if "sErrorCode" in body or "AADSTS" in body:
            _extract_azure_error(body, config.correlation_id)
        return url, body

    # ------------------------------------------------------------------
    # Step 4b helper: follow oPostParams
    # ------------------------------------------------------------------

    def _follow_o_post_params(self, html: str, correlation_id: str) -> tuple[str, str]:
        v = _extract_dollar_config(html)

        url_post_raw = v.get("urlPost", "")
        if not url_post_raw:
            raise AzureFormParsingError("oPostParams page: missing urlPost in $Config")
        url_post = (
            url_post_raw
            if url_post_raw.startswith("http")
            else f"{_AZURE_BASE}{url_post_raw}"
        )

        o_post_params = v.get("oPostParams")
        if not isinstance(o_post_params, dict):
            raise AzureFormParsingError(
                "oPostParams page: missing or non-dict oPostParams in $Config"
            )

        form = {k: str(val) for k, val in o_post_params.items()}
        logger.info("Following oPostParams to %s", url_post)

        url, body = self._client.post_form_following_redirects(url_post, form)
        if "AADSTS" in body or "sErrorCode" in body:
            _extract_azure_error(body, correlation_id)
        return url, body

    # ------------------------------------------------------------------
    # Step 5: ACS submission
    # ------------------------------------------------------------------

    def _process_acs_response(self, url: str, html: str) -> AuthSession:
        saml_response = _extract_form_field(html, "SAMLResponse")
        if not saml_response:
            raise SamlAssertionError("SAMLResponse field not found in ACS form")

        relay_state = _extract_form_field(html, "RelayState") or ""
        form_action = (
            _extract_form_action(html)
            or f"{self._config.base_url}/authentication/saml/acs"
        )

        logger.debug("Submitting SAMLResponse to %s", form_action)
        form = {"SAMLResponse": saml_response, "RelayState": relay_state}
        final_url, final_body = self._client.post_form_following_redirects(
            form_action, form
        )
        logger.debug("ACS response URL: %s", final_url)

        return self._build_session_from_saml_response(
            saml_response, final_url, final_body
        )

    # ------------------------------------------------------------------
    # Step 6: Parse SAML assertion and build session
    # ------------------------------------------------------------------

    def _build_session_from_saml_response(
        self, saml_response_b64: str, _final_url: str, _final_body: str
    ) -> AuthSession:
        saml = parse_saml_response_base64(saml_response_b64)
        validate_response(
            saml,
            self._config.sp_acs_url,
            clock_skew_seconds=int(self._config.clock_skew_tolerance.total_seconds()),
        )

        if saml.assertion is None:
            raise SamlAssertionError("No assertion in SAML response")

        user_info = saml.assertion.to_user_info()

        expires_at = (
            saml.assertion.session_not_on_or_after
            or saml.assertion.not_on_or_after
            or (datetime.now(timezone.utc) + timedelta(hours=1))
        )

        session = AuthSession(user_info=user_info, expires_at=expires_at)
        logger.info("Authentication successful for %s", session.user_info.username)
        return session
