"""MTA channel implementation for handling email delivery."""

import hashlib
import logging
import secrets

from django.conf import settings

import jwt
from drf_spectacular.utils import extend_schema
from jmap_email import parse_email
from rest_framework import status, viewsets
from rest_framework.authentication import BaseAuthentication
from rest_framework.decorators import action
from rest_framework.exceptions import AuthenticationFailed
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response

from core import enums, models
from core.mda.inbound import check_local_recipients, deliver_inbound_message
from core.mda.raw_mime import remove_mime_headers

logger = logging.getLogger(__name__)


class MTAJWTAuthentication(BaseAuthentication):
    """Authenticate the MTA-to-MDA channel via an HS256 JWT.

    Trust model: the whole channel rests on the shared ``MDA_API_SECRET``.
    Only the MTA-in service knows it, so a valid HMAC signature *is* the proof
    of identity — there is no per-request identity beyond "signed by the
    secret". Consequently we do NOT attempt replay protection (a ``jti`` nonce
    store, etc.): anyone able to forge the signature already holds the secret
    and could mint fresh tokens at will, and anyone who cannot is stopped by
    the signature check. Keeping the secret out of source (it has no default —
    see ``settings.MDA_API_SECRET``) and the transport on TLS is what actually
    secures this path.

    On top of the signature we keep two cheap, narrow guards:
    - ``exp``: bounds a leaked token's useful lifetime. The issuer sizes the
      claim to cover its full retry window (see mta-in ``mda_api_call``).
    - ``body_hash``: binds the token to its exact request body, so a captured
      token can't be repurposed for a *different* body within that window.
      Enforced even for an empty body (the bodyless ``/check`` path).

    Returns None or (user, auth).
    """

    def authenticate(self, request):
        auth_header = request.headers.get("Authorization")
        if not auth_header:
            return None

        try:
            jwt_token = auth_header.split(" ")[1]
            payload = jwt.decode(
                jwt_token,
                settings.MDA_API_SECRET,
                algorithms=["HS256"],
                options={
                    # exp bounds the lifetime; body_hash binds the token to its
                    # payload. Both are mandatory.
                    "require": ["exp", "body_hash"],
                    "verify_exp": True,
                    "verify_signature": True,
                },
            )

            # Bind the token to its payload. Always enforced — including for
            # an empty body (sha256 of b"") — so the bodyless /check endpoint
            # can't be driven with a token minted for a different request.
            claimed_hash = payload["body_hash"]
            # ``compare_digest`` raises TypeError on mismatched types (e.g. a
            # numeric ``body_hash`` claim), which would surface as a 500 rather
            # than an auth failure. Reject a non-string claim up front.
            if not isinstance(claimed_hash, str):
                raise jwt.InvalidTokenError("Invalid email hash")
            body_hash = hashlib.sha256(request.body or b"").hexdigest()
            if not secrets.compare_digest(body_hash, claimed_hash):
                raise jwt.InvalidTokenError("Invalid email hash")

            service_account = models.User()
            return (service_account, payload)

        except (jwt.ExpiredSignatureError, jwt.InvalidTokenError) as e:
            raise AuthenticationFailed("Invalid token") from e
        except (IndexError, KeyError) as e:
            raise AuthenticationFailed("Invalid token header or payload") from e

    def authenticate_header(self, request):
        """Return the header to be used in the WWW-Authenticate response header."""
        return 'Bearer realm="MTA"'


class InboundMTAViewSet(viewsets.GenericViewSet):
    """Handles incoming email messages from MTA (Mail Transfer Agent)."""

    # Channel metadata
    CHANNEL_TYPE = "mta"
    CHANNEL_DESCRIPTION = "Mail Transfer Agent (email)"

    permission_classes = [IsAuthenticated]
    authentication_classes = [MTAJWTAuthentication]

    @extend_schema(exclude=True)
    @action(
        detail=False, methods=["post"], url_path="check", url_name="inbound-mta-check"
    )
    def check(self, request):
        """Check recipients exist."""
        data = request.data
        addresses = data.get("addresses", [])
        if not addresses or not isinstance(addresses, list):
            return Response(
                {"detail": "Missing addresses"}, status=status.HTTP_400_BAD_REQUEST
            )

        local_addresses = check_local_recipients(addresses)
        results = {address: address in local_addresses for address in addresses}
        return Response(results)

    @extend_schema(exclude=True)
    @action(
        detail=False,
        methods=["post"],
        url_path="deliver",
        url_name="inbound-mta-deliver",
    )
    def deliver(self, request):
        """Handle incoming raw email (message/rfc822) from MTA."""

        # request.user will be the service account, request.auth the JWT payload
        mta_metadata = request.auth
        if not mta_metadata or "original_recipients" not in mta_metadata:
            # This case should ideally be caught by the authentication class
            logger.error("MTA metadata missing or malformed in authenticated request.")
            return Response(
                {"detail": "Internal authentication error"},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )

        # Validate content type (optional but good practice)
        # Note: If parser_classes included FormParser or MultiPartParser, request.body might be consumed.
        # Ensure parser_classes=[parsers.BaseParser] or similar if relying on request.body.
        if request.content_type != "message/rfc822":
            logger.warning(
                "Received inbound POST with incorrect Content-Type: %s",
                request.content_type,
            )
            # Decide whether to reject or attempt parsing anyway
            return Response(
                {"detail": "Content-Type must be message/rfc822"},
                status=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
            )

        raw_data = request.body
        if not raw_data:
            logger.error("Received empty body for inbound email.")
            return Response(
                {"status": "error", "detail": "Empty request body"},
                status=status.HTTP_400_BAD_REQUEST,
            )

        # Validate incoming email size
        email_size = len(raw_data)
        if email_size > settings.MAX_INCOMING_EMAIL_SIZE:
            logger.warning(
                "Incoming email size (%d bytes) exceeds maximum allowed size (%d bytes)",
                email_size,
                settings.MAX_INCOMING_EMAIL_SIZE,
            )
            return Response(
                {
                    "status": "error",
                    "detail": f"Incoming email size ({email_size} bytes) exceeds maximum allowed size "
                    + f"({settings.MAX_INCOMING_EMAIL_SIZE} bytes)",
                },
                status=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
            )

        logger.info(
            "Raw email received: %d bytes for %s",
            len(raw_data),
            mta_metadata["original_recipients"],  # Log all intended recipients
        )

        # Strip sender-supplied headers we own or authoritatively rewrite:
        #  - ``X-StMsg-*``: our pipeline namespace (sender-auth etc.) — must not
        #    be forgeable.
        #  - ``Return-Path``: only the delivering MTA may write it (RFC 5321
        #    §4.4); we rewrite it below from the real envelope MAIL FROM, so any
        #    inbound copy is forged and must go first.
        raw_data = remove_mime_headers(
            raw_data, prefixes=["x-stmsg-"], names=["return-path"]
        )

        def sanitize_header(header: str) -> str:
            return header.replace("\r", "").replace("\n", "")[0:255]

        # Bake the immutable envelope facts as standard headers at ingest, on
        # top of the received bytes and BEFORE the blob is created, so they ride
        # in the single stored blob:
        #  - ``Return-Path``: the envelope MAIL FROM (``<>`` for a null sender /
        #    bounce). Durable home for what the autoreply/bounce logic reads.
        #  - ``Received``: the SMTP trace (HELO / rDNS / IP), when the MTA
        #    forwarded it. Shared across recipients (one prepend), so identical
        #    bytes still dedup to one blob.
        sender = mta_metadata.get("sender") or ""
        prepend_headers = [("Return-Path", f"<{sender}>" if sender else "<>")]
        client_helo = mta_metadata.get("client_helo")
        client_hostname = mta_metadata.get("client_hostname")
        client_address = mta_metadata.get("client_address")

        # Only emit the Received trace when the MTA forwarded all three parts
        if client_helo and client_hostname and client_address:
            prepend_headers.append(
                (
                    "Received",
                    f"from {client_helo} ({client_hostname} [{client_address}]);",
                )
            )

        raw_data = (
            "\r\n".join([f"{k}: {sanitize_header(v)}" for k, v in prepend_headers])
            + "\r\n"
        ).encode("utf-8") + raw_data

        # Parse the email message once
        parsed_email = parse_email(raw_data)
        if parsed_email is None:
            # Sender-supplied malformed input; not an internal error.
            logger.warning("Failed to parse inbound email (returning 400)")
            # Consider saving the raw email for debugging
            return Response(
                {"status": "error", "detail": "Failed to parse email"},
                status=status.HTTP_400_BAD_REQUEST,  # Bad request as email is malformed
            )

        # Structured SMTP envelope carried alongside the bytes (never injected
        # into them, so the blob stays byte-identical across recipients and
        # dedups to one row). ``sender`` is the MAIL FROM; the client_* fields
        # are the connecting SMTP peer (forwarded by the MTA). ``rcpt_to`` is
        # the actual RCPT TO and is set per recipient inside the loop.
        base_envelope = {
            "origin": enums.InboundOrigin.MTA,
            "mail_from": sender,
            "ip": client_address or "",
            "helo": client_helo or "",
            "hostname": client_hostname or "",
        }

        # Deliver the parsed email to each original recipient
        success_count = 0
        failure_count = 0
        delivery_results = {}

        for recipient in mta_metadata["original_recipients"]:
            try:
                # Call the refactored delivery function which returns True/False
                delivered = deliver_inbound_message(
                    recipient,
                    parsed_email,
                    raw_data,
                    envelope={**base_envelope, "rcpt_to": recipient},
                )
                if delivered:
                    success_count += 1
                    delivery_results[recipient] = "Success"
                else:
                    # Delivery function failed (and logged the reason)
                    failure_count += 1
                    delivery_results[recipient] = "Failed"
            except Exception as e:  # pylint: disable=broad-exception-caught
                logger.error(
                    "Unexpected error during delivery loop for %s: %s",
                    recipient,
                    e,
                    exc_info=True,
                )
                failure_count += 1
                delivery_results[recipient] = f"Error: {e}"

        # Determine overall status based on counts
        if failure_count > 0 and success_count == 0:
            # If all deliveries failed, return a server error
            logger.error("All deliveries failed for inbound email.")
            return Response(
                {
                    "status": "error",
                    "detail": "Failed to deliver message to any recipient",
                    "results": delivery_results,
                },
                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
            )

        if failure_count > 0:
            # If some deliveries failed, return 207 Multi-Status
            logger.warning(
                "Partial delivery failure: %d successful, %d failed",
                success_count,
                failure_count,
            )
            return Response(
                {
                    "status": "partial_success",
                    "delivered": success_count,
                    "failed": failure_count,
                    "results": delivery_results,
                },
                status=status.HTTP_207_MULTI_STATUS,
            )

        # All deliveries successful
        logger.info("All %d deliveries successful for inbound email.", success_count)
        return Response({"status": "ok", "delivered": success_count})
