#!/usr/bin/env python3
"""Secret-holding Orga mail-draft service, scoped to the account's Drafts folder.

Writes go through IMAP APPEND (create), APPEND + UID STORE/UID EXPUNGE
(update: new message first, then remove the old one) and UID STORE/UID
EXPUNGE (delete). UIDPLUS is mandatory: the service never falls back to a
mailbox-wide EXPUNGE. The service resolves the Drafts folder itself via
IMAP LIST (preferring the RFC 6154 \\Drafts special-use flag, falling back to
a small set of known names) and never accepts a folder name from the caller,
so it structurally cannot touch any other mailbox. No SMTP access, no
mail-send: putting a message in Drafts is as far as this goes (see
hosts/blunix/localServices/orga-agent/README.md, "Deliberately unavailable:
mail-send").

Every draft this service creates carries an `X-Orga-Draft: 1` marker header;
update/delete/show/list only ever act on messages carrying that marker, so a
draft the human started composing themselves in their own mail client is
never touched. The caller-facing "id" is the message's own Message-ID
(generated by this service, format `<hex@orga-mail-draft.internal>`) instead
of the old local UUID -- it stays stable across `update` calls the same way
the UUID used to.
"""

import datetime as dt
import email
import email.policy
import email.utils
import hashlib
import imaplib
import json
import os
import re
import socket
import sys
import uuid
from email.message import EmailMessage
from pathlib import Path

MAX_REQUEST = 65_536
CONTROL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]")
HEADER_UNSAFE = re.compile(r"[\r\n]")
ADDRESS = re.compile(r"[^\s@<>,;]{1,64}@[^\s@<>,;]{1,255}\.[^\s@<>,;]{2,24}\Z")
DRAFT_ID = re.compile(r"<[0-9a-f]{32}@orga-mail-draft\.internal>\Z")
CONTENT_FIELDS = ("to", "cc", "bcc", "subject", "body", "attachments")

# RFC 6154 SPECIAL-USE is the primary way to find the Drafts folder --
# language- and encoding-agnostic. These ASCII/German names are only a
# fallback for a server that doesn't advertise SPECIAL-USE.
FOLDER_LINE = re.compile(rb'^\((?P<flags>[^)]*)\)\s+(?:"(?P<delim>[^"]*)"|NIL)\s+(?P<name>.+)$')
FALLBACK_DRAFTS_NAMES = ("Drafts", "Draft", "Entwürfe", "INBOX.Drafts", "INBOX/Drafts", "INBOX.Entwürfe")

# SEARCH HEADER on a custom X-header is not reliably indexed by every real
# IMAP server (confirmed against live Zoho: SEARCH HEADER X-Orga-Draft 1
# returned no matches for a message that demonstrably had that header).
# SINCE this fixed epoch is a broad, standards-supported stand-in for "all
# messages in the folder"; ownership is then verified client-side.
EPOCH_SEARCH_DATE = "01-Jan-1970"


class RequestError(Exception):
    pass


def secret(path_variable):
    value = Path(os.environ[path_variable]).read_text(encoding="utf-8").strip()
    if not value:
        raise RuntimeError("configured credential is empty")
    return value


def text(value, name, maximum, allow_empty=False):
    if not isinstance(value, str) or len(value) > maximum or CONTROL.search(value):
        raise RequestError(f"invalid {name}")
    if not allow_empty and not value.strip():
        raise RequestError(f"missing {name}")
    return value


def header_text(value, name, maximum):
    value = text(value, name, maximum)
    if HEADER_UNSAFE.search(value):
        raise RequestError(f"invalid {name}")
    return value


def address_list(value, name, minimum):
    if not isinstance(value, list) or len(value) > 50:
        raise RequestError(f"invalid {name}")
    result = [text(item, name, 320) for item in value]
    for item in result:
        if not ADDRESS.fullmatch(item):
            raise RequestError(f"invalid address in {name}: {item!r}")
    if len(result) < minimum:
        raise RequestError(f"missing {name}")
    return result


def attachment_list(value):
    if not isinstance(value, list) or len(value) > 50:
        raise RequestError("invalid attachments")
    return [header_text(item, "attachments", 300) for item in value]


def draft_id(value):
    if not isinstance(value, str) or not DRAFT_ID.fullmatch(value):
        raise RequestError("invalid id")
    return value


def content_hash(content):
    canonical = json.dumps(content, sort_keys=True, ensure_ascii=False)
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def parsed_content(request):
    if set(request) - {"operation", "id", *CONTENT_FIELDS}:
        raise RequestError("unsupported draft request")
    content = {
        "to": address_list(request.get("to"), "to", minimum=1),
        "cc": address_list(request.get("cc", []), "cc", minimum=0),
        "bcc": address_list(request.get("bcc", []), "bcc", minimum=0),
        "subject": header_text(request.get("subject"), "subject", 500),
        "body": text(request.get("body"), "body", 50_000),
        "attachments": attachment_list(request.get("attachments", [])),
    }
    return content


def mail_login():
    connection = imaplib.IMAP4_SSL(os.environ["ORGA_MAIL_HOST"], timeout=20)
    connection.login(os.environ["ORGA_MAIL_USER"], secret("ORGA_MAIL_PASSWORD_FILE"))
    return connection


def discover_drafts_folder(connection):
    status, values = connection.list()
    if status != "OK":
        raise RuntimeError("IMAP LIST failed while locating the drafts folder")
    candidates = []
    for raw in values or []:
        line = raw if isinstance(raw, bytes) else raw[0]
        match = FOLDER_LINE.match(line)
        if not match:
            continue
        wire_name = match.group("name").decode("utf-8", errors="replace").strip()
        display_name = wire_name[1:-1] if wire_name.startswith('"') and wire_name.endswith('"') else wire_name
        if any(flag.lower() == b"\\drafts" for flag in match.group("flags").split()):
            return wire_name
        candidates.append((display_name, wire_name))
    for fallback in FALLBACK_DRAFTS_NAMES:
        for display_name, wire_name in candidates:
            if display_name == fallback:
                return wire_name
    raise RuntimeError(
        "could not identify the drafts folder via IMAP LIST "
        "(no \\Drafts special-use flag and no recognized fallback name)"
    )


def select_folder(connection, folder, readonly):
    status, _ = connection.select(folder, readonly=readonly)
    if status != "OK":
        raise RuntimeError(f"could not select the drafts folder {folder!r}")


def new_message_id():
    return f"<{uuid.uuid4().hex}@orga-mail-draft.internal>"


def build_message(content, message_id, revision, created_at, updated_at):
    digest = content_hash(content)
    message = EmailMessage()
    message["From"] = os.environ["ORGA_MAIL_USER"]
    message["To"] = ", ".join(content["to"])
    if content["cc"]:
        message["Cc"] = ", ".join(content["cc"])
    if content["bcc"]:
        message["Bcc"] = ", ".join(content["bcc"])
    message["Subject"] = content["subject"]
    message["Date"] = email.utils.formatdate(localtime=False)
    message["Message-ID"] = message_id
    message["X-Orga-Draft"] = "1"
    message["X-Orga-Draft-Revision"] = str(revision)
    message["X-Orga-Draft-Created"] = created_at
    message["X-Orga-Draft-Updated"] = updated_at
    message["X-Orga-Draft-Content-Hash"] = digest
    for attachment in content["attachments"]:
        message["X-Orga-Draft-Attachment"] = attachment
    message.set_content(content["body"])
    return message.as_bytes(), digest


def append_message(connection, folder, raw):
    status, _ = connection.append(folder, "(\\Draft)", None, raw)
    if status != "OK":
        raise RuntimeError("IMAP APPEND failed while writing the draft")


def uid_sort_key(uid_bytes):
    try:
        return int(uid_bytes)
    except (TypeError, ValueError):
        return -1


def revision(message):
    try:
        value = int(message.get("X-Orga-Draft-Revision", ""))
    except ValueError:
        return None
    return value if value >= 1 else None


def owned_matches(connection, the_id, fetch_spec="(BODY.PEEK[])"):
    status, values = connection.uid("search", None, "HEADER", "Message-ID", the_id)
    if status != "OK":
        raise RuntimeError("IMAP UID SEARCH failed")
    matches = []
    for uid_bytes in (values[0] or b"").split():
        status, data = connection.uid("fetch", uid_bytes, fetch_spec)
        if status != "OK" or not data or not isinstance(data[0], tuple):
            continue
        message = email.message_from_bytes(data[0][1], policy=email.policy.default)
        if (
            message.get("X-Orga-Draft") == "1"
            and str(message.get("Message-ID", "")).strip() == the_id
            and revision(message) is not None
        ):
            matches.append((uid_bytes, message))
    if not matches:
        raise RequestError("draft was not found")
    # Revision is monotonically increasing for this service. The timestamp
    # and UID make recovery deterministic even for malformed server clocks.
    matches.sort(
        key=lambda item: (
            revision(item[1]),
            item[1].get("X-Orga-Draft-Updated") or "",
            uid_sort_key(item[0]),
        ),
        reverse=True,
    )
    return matches


def load_owned_message(connection, the_id):
    return owned_matches(connection, the_id)[0]


def require_uidplus(connection):
    status, values = connection.capability()
    if status != "OK" or not any(b"UIDPLUS" in value.upper().split() for value in values or []):
        raise RuntimeError("IMAP UIDPLUS is required for safe draft deletion")


def delete_uid(connection, uid_bytes):
    require_uidplus(connection)
    status, _ = connection.uid("STORE", uid_bytes, "+FLAGS", "(\\Deleted)")
    if status != "OK":
        raise RuntimeError("IMAP STORE failed while removing the previous draft")
    status, _ = connection.uid("EXPUNGE", uid_bytes)
    if status != "OK":
        raise RuntimeError("IMAP UID EXPUNGE failed while removing the previous draft")


def split_addresses(header_value):
    if header_value is None:
        return []
    return [address.addr_spec for address in header_value.addresses]


def extract_body(message):
    if message.is_multipart():
        return ""
    content = message.get_content()
    # set_content() always appends exactly one trailing newline to a
    # text/plain payload; undo that one so a round trip returns exactly
    # what was originally submitted, whatever its own trailing whitespace.
    return content[:-1] if content.endswith("\n") else content


def message_to_draft(the_id, message):
    return {
        "id": the_id,
        "revision": revision(message),
        "created_at": message.get("X-Orga-Draft-Created") or "",
        "updated_at": message.get("X-Orga-Draft-Updated") or "",
        "content_hash": message.get("X-Orga-Draft-Content-Hash", ""),
        "to": split_addresses(message.get("To")),
        "cc": split_addresses(message.get("Cc")),
        "bcc": split_addresses(message.get("Bcc")),
        "subject": str(message.get("Subject", "")),
        "body": extract_body(message),
        "attachments": [str(value) for value in message.get_all("X-Orga-Draft-Attachment", [])],
    }


def create_draft(connection, folder, content):
    now = dt.datetime.now(dt.timezone.utc).isoformat()
    message_id = new_message_id()
    raw, digest = build_message(content, message_id, revision=1, created_at=now, updated_at=now)
    append_message(connection, folder, raw)
    return {
        "id": message_id, "revision": 1,
        "created_at": now, "updated_at": now,
        "content_hash": digest, **content,
    }


def update_draft(connection, folder, the_id, content):
    # Look up and validate the old message -- and remember its UID -- before
    # appending the replacement, which reuses the same Message-ID. Doing it
    # in this order (append first, delete old second) means a failure after
    # the append still leaves a valid draft (a harmless duplicate to clean
    # up), instead of a failure after a delete-first step losing the draft
    # outright.
    old_uid, old_message = load_owned_message(connection, the_id)
    # Verify the safety prerequisite before APPEND. Otherwise a server that
    # lacks UIDPLUS could leave a new duplicate although the update must fail.
    require_uidplus(connection)
    next_revision = revision(old_message) + 1
    created_at = old_message.get("X-Orga-Draft-Created") or dt.datetime.now(dt.timezone.utc).isoformat()
    updated_at = dt.datetime.now(dt.timezone.utc).isoformat()
    raw, digest = build_message(content, the_id, next_revision, created_at, updated_at)
    append_message(connection, folder, raw)
    delete_uid(connection, old_uid)
    # A previous partial update can have left older, valid own revisions.
    # Clean them up only after the new revision is safely present, and only
    # through the same UID-bound deletion primitive.
    for stale_uid, stale_message in owned_matches(connection, the_id):
        if revision(stale_message) < next_revision:
            delete_uid(connection, stale_uid)
    return {
        "id": the_id, "revision": next_revision,
        "created_at": created_at, "updated_at": updated_at,
        "content_hash": digest, **content,
    }


def delete_draft(connection, the_id):
    # Every deletion is individually UID-bound; unrelated \Deleted messages
    # in the selected folder cannot be expunged by this operation.
    for uid_bytes, _ in owned_matches(connection, the_id):
        delete_uid(connection, uid_bytes)


def list_drafts(connection):
    status, values = connection.uid("search", None, "SINCE", EPOCH_SEARCH_DATE)
    if status != "OK":
        raise RuntimeError("IMAP UID SEARCH failed")
    newest = {}
    for uid_bytes in (values[0] or b"").split():
        status, data = connection.uid("fetch", uid_bytes, "(BODY.PEEK[HEADER])")
        if status != "OK" or not data or not isinstance(data[0], tuple):
            continue
        message = email.message_from_bytes(data[0][1], policy=email.policy.default)
        if message.get("X-Orga-Draft") != "1":
            continue
        message_id = str(message.get("Message-ID", "")).strip()
        if not DRAFT_ID.fullmatch(message_id) or revision(message) is None:
            continue
        candidate = (revision(message), message.get("X-Orga-Draft-Updated") or "", uid_sort_key(uid_bytes))
        previous = newest.get(message_id)
        if previous is not None and candidate <= previous[0]:
            continue
        newest[message_id] = (candidate, {
            "id": message_id,
            "revision": revision(message),
            "subject": str(message.get("Subject", "")),
            "updated_at": message.get("X-Orga-Draft-Updated") or "",
        })
    drafts = [item[1] for item in newest.values()]
    drafts.sort(key=lambda draft: draft["updated_at"])
    return drafts


def draft_request(request):
    operation = request.get("operation")
    if operation == "capabilities" and set(request) == {"operation"}:
        return {
            "ok": True, "capability": "mail-draft", "writes": "imap-drafts-folder-only",
            "network": True, "operations": ["create", "update", "show", "list", "delete"],
        }

    content = None
    the_id = None
    if operation == "list" and set(request) == {"operation"}:
        pass
    elif operation == "create":
        content = parsed_content(request)
    elif operation == "update":
        if "id" not in request:
            raise RequestError("missing id")
        the_id = draft_id(request["id"])
        content = parsed_content(request)
    elif operation == "show" and set(request) == {"operation", "id"}:
        the_id = draft_id(request["id"])
    elif operation == "delete" and set(request) == {"operation", "id"}:
        the_id = draft_id(request["id"])
    else:
        raise RequestError("unsupported draft request")

    connection = mail_login()
    try:
        folder = discover_drafts_folder(connection)

        if operation == "list":
            select_folder(connection, folder, readonly=True)
            return {"ok": True, "trust": "UNTRUSTED_EXTERNAL_DATA", "drafts": list_drafts(connection)}

        if operation == "create":
            select_folder(connection, folder, readonly=False)
            return {"ok": True, "operation": "create", "draft": create_draft(connection, folder, content)}

        if operation == "update":
            select_folder(connection, folder, readonly=False)
            return {"ok": True, "operation": "update", "draft": update_draft(connection, folder, the_id, content)}

        if operation == "show":
            select_folder(connection, folder, readonly=True)
            _, message = load_owned_message(connection, the_id)
            return {
                "ok": True,
                "trust": "UNTRUSTED_EXTERNAL_DATA",
                "draft": message_to_draft(the_id, message),
            }

        if operation == "delete":
            select_folder(connection, folder, readonly=False)
            delete_draft(connection, the_id)
            return {"ok": True, "operation": "delete", "id": the_id}
    finally:
        try:
            connection.logout()
        except imaplib.IMAP4.error:
            pass


def serve():
    listener = socket.socket(fileno=0)
    while True:
        connection, _ = listener.accept()
        with connection:
            try:
                raw = b""
                while not raw.endswith(b"\n") and len(raw) <= MAX_REQUEST:
                    chunk = connection.recv(4096)
                    if not chunk:
                        break
                    raw += chunk
                if len(raw) > MAX_REQUEST:
                    raise RequestError("request is too large")
                request = json.loads(raw)
                if not isinstance(request, dict):
                    raise RequestError("request must be an object")
                response = draft_request(request)
            except (RequestError, json.JSONDecodeError) as error:
                response = {"ok": False, "error": str(error)}
            except Exception:
                response = {"ok": False, "error": "upstream service request failed"}
            connection.sendall(json.dumps(response, ensure_ascii=False).encode("utf-8"))


def main():
    serve()


if __name__ == "__main__":
    sys.exit(main() or 0)