"""SAML 2.0 response parser.

Parses the base64-encoded SAMLResponse posted to the ACS endpoint.
Uses lxml for robust XML handling.
"""

from __future__ import annotations

import base64
from datetime import datetime, timezone
from typing import Optional
from xml.etree import ElementTree as ET

from .errors import SamlAssertionError, SamlParseError, SamlValidationError
from .types import SamlAssertion, SamlResponse, SamlStatus

# XML namespace map
_NS = {
    "samlp": "urn:oasis:names:tc:SAML:2.0:protocol",
    "saml": "urn:oasis:names:tc:SAML:2.0:assertion",
}


def _tag(ns: str, local: str) -> str:
    """Build a Clark-notation tag like {urn:...}LocalName."""
    return f"{{{_NS[ns]}}}{local}"


def _parse_dt(value: Optional[str]) -> Optional[datetime]:
    if not value:
        return None
    # Python 3.11+ handles 'Z' directly; for older versions replace manually
    value = value.replace("Z", "+00:00")
    try:
        return datetime.fromisoformat(value)
    except ValueError:
        return None


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def parse_saml_response_base64(encoded: str) -> SamlResponse:
    """Decode and parse a base64-encoded SAMLResponse."""
    try:
        xml_bytes = base64.b64decode(encoded)
    except Exception as exc:
        raise SamlParseError(f"Failed to base64-decode SAMLResponse: {exc}") from exc

    try:
        xml_text = xml_bytes.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise SamlParseError(f"SAMLResponse is not valid UTF-8: {exc}") from exc

    return parse_saml_response(xml_text)


def parse_saml_response(xml_text: str) -> SamlResponse:
    """Parse a SAMLResponse XML string into a SamlResponse object."""
    try:
        root = ET.fromstring(xml_text)
    except ET.ParseError as exc:
        raise SamlParseError(
            f"XML parsing error: {exc}", xml_snippet=xml_text[:200]
        ) from exc

    # ---- Response attributes ----
    response_id = root.get("ID", "")
    if not response_id:
        raise SamlParseError(
            "SAML Response missing ID attribute",
            xml_snippet=xml_text[:200],
        )

    in_response_to = root.get("InResponseTo")
    destination = root.get("Destination")
    issue_instant = _parse_dt(root.get("IssueInstant")) or datetime.now(timezone.utc)

    # ---- Issuer ----
    issuer_el = root.find(_tag("saml", "Issuer"))
    issuer = (issuer_el.text or "").strip() if issuer_el is not None else ""

    # ---- Status ----
    status_el = root.find(f"{_tag('samlp', 'Status')}/{_tag('samlp', 'StatusCode')}")
    if status_el is None:
        raise SamlParseError(
            "SAML Response missing StatusCode",
            xml_snippet=xml_text[:200],
        )
    status_code = status_el.get("Value", "")

    status_msg_el = root.find(
        f"{_tag('samlp', 'Status')}/{_tag('samlp', 'StatusMessage')}"
    )
    status_message = (
        (status_msg_el.text or "").strip() if status_msg_el is not None else None
    )

    status = SamlStatus(
        status_code=status_code,
        status_message=status_message or None,
    )

    if not status.is_success():
        raise SamlValidationError(
            f"SAML authentication failed with status: {status_code}",
            response_status=status_code,
        )

    # ---- Assertion ----
    assertion_el = root.find(_tag("saml", "Assertion"))
    assertion = _parse_assertion(assertion_el) if assertion_el is not None else None

    return SamlResponse(
        id=response_id,
        in_response_to=in_response_to,
        issue_instant=issue_instant,
        status=status,
        issuer=issuer,
        assertion=assertion,
        destination=destination,
    )


def _parse_assertion(el: ET.Element) -> SamlAssertion:
    assertion_id = el.get("ID", "")
    issue_instant = _parse_dt(el.get("IssueInstant")) or datetime.now(timezone.utc)

    issuer_el = el.find(_tag("saml", "Issuer"))
    issuer = (issuer_el.text or "").strip() if issuer_el is not None else ""

    # Subject / NameID
    name_id_el = el.find(f"{_tag('saml', 'Subject')}/{_tag('saml', 'NameID')}")
    subject = ""
    name_id_format: Optional[str] = None
    if name_id_el is not None:
        subject = (name_id_el.text or "").strip()
        name_id_format = name_id_el.get("Format")

    # Conditions
    conditions_el = el.find(_tag("saml", "Conditions"))
    not_before = None
    not_on_or_after = None
    audience_restrictions: list[str] = []
    if conditions_el is not None:
        not_before = _parse_dt(conditions_el.get("NotBefore"))
        not_on_or_after = _parse_dt(conditions_el.get("NotOnOrAfter"))
        for aud_el in conditions_el.findall(
            f"{_tag('saml', 'AudienceRestriction')}/{_tag('saml', 'Audience')}"
        ):
            if aud_el.text:
                audience_restrictions.append(aud_el.text.strip())

    # AuthnStatement
    authn_el = el.find(_tag("saml", "AuthnStatement"))
    session_index: Optional[str] = None
    session_not_on_or_after: Optional[datetime] = None
    authn_instant: Optional[datetime] = None
    authn_context_class_ref: Optional[str] = None
    if authn_el is not None:
        session_index = authn_el.get("SessionIndex")
        session_not_on_or_after = _parse_dt(authn_el.get("SessionNotOnOrAfter"))
        authn_instant = _parse_dt(authn_el.get("AuthnInstant"))
        ctx_el = authn_el.find(
            f"{_tag('saml', 'AuthnContext')}/{_tag('saml', 'AuthnContextClassRef')}"
        )
        if ctx_el is not None:
            authn_context_class_ref = (ctx_el.text or "").strip() or None

    # Attributes
    attributes: dict[str, list[str]] = {}
    attr_stmt = el.find(_tag("saml", "AttributeStatement"))
    if attr_stmt is not None:
        for attr_el in attr_stmt.findall(_tag("saml", "Attribute")):
            name = attr_el.get("Name", "")
            if not name:
                continue
            values: list[str] = []
            for val_el in attr_el.findall(_tag("saml", "AttributeValue")):
                if val_el.text:
                    values.append(val_el.text.strip())
            attributes[name] = values

    return SamlAssertion(
        id=assertion_id,
        issue_instant=issue_instant,
        issuer=issuer,
        subject=subject,
        name_id_format=name_id_format,
        not_before=not_before,
        not_on_or_after=not_on_or_after,
        session_index=session_index,
        session_not_on_or_after=session_not_on_or_after,
        authn_context_class_ref=authn_context_class_ref,
        authn_instant=authn_instant,
        attributes=attributes,
        audience_restrictions=audience_restrictions,
        signature_verified=False,
    )


# ---------------------------------------------------------------------------
# Validation
# ---------------------------------------------------------------------------


def validate_response(
    response: SamlResponse,
    expected_acs_url: str,
    clock_skew_seconds: int = 300,
) -> None:
    """Validate a parsed SamlResponse.  Raises SamlValidationError on failure."""
    from datetime import timedelta

    if not response.status.is_success():
        raise SamlValidationError(
            f"SAML Response failed: {response.status.status_code}",
            response_status=response.status.status_code,
        )

    if response.destination and response.destination != expected_acs_url:
        raise SamlValidationError(
            f"SAML Response destination mismatch: "
            f"expected {expected_acs_url!r}, got {response.destination!r}"
        )

    if response.assertion is not None:
        skew = timedelta(seconds=clock_skew_seconds)
        if not response.assertion.is_valid(clock_skew=skew):
            raise SamlValidationError("SAML Assertion has expired or is not yet valid")


def extract_user_info_from_response(response: SamlResponse):
    """Convenience wrapper: extract UserInfo from a parsed SamlResponse."""
    if response.assertion is None:
        raise SamlAssertionError("SAML Response contains no assertion")
    return response.assertion.to_user_info()
