from dataclasses import dataclass, field
from email import policy
from email.message import EmailMessage, Message
from email.parser import BytesParser


@dataclass
class ExtractedMimePart:
    filename: str
    content_type: str
    payload: bytes
    content_id: str | None = None
    disposition: str = "attachment"


@dataclass
class MirroredMimeContent:
    body_text: str = ""
    body_html: str | None = None
    inline_parts: list[ExtractedMimePart] = field(default_factory=list)
    file_attachments: list[ExtractedMimePart] = field(default_factory=list)


def _part_payload(part: Message) -> bytes:
    payload = part.get_payload(decode=True)
    if isinstance(payload, bytes):
        return payload
    if isinstance(payload, str):
        return payload.encode("utf-8", errors="replace")
    return b""


def _part_text(part: Message) -> str:
    try:
        value = part.get_content()
    except (LookupError, UnicodeError):
        value = _part_payload(part).decode("utf-8", errors="replace")
    return value if isinstance(value, str) else ""


def _content_id(part: Message) -> str | None:
    header = part.get("Content-ID")
    if not header:
        return None
    return header.strip().strip("<>")


def _is_attachment_part(part: Message) -> bool:
    disposition = (part.get_content_disposition() or "").lower()
    if disposition == "attachment":
        return True
    if disposition == "inline":
        return False
    filename = part.get_filename()
    content_type = part.get_content_type()
    if filename and content_type not in {"text/plain", "text/html"}:
        return True
    return False


def _walk_parts(message: Message, collector: MirroredMimeContent) -> None:
    if message.is_multipart():
        for part in message.get_payload():
            if isinstance(part, Message):
                _walk_parts(part, collector)
        return

    content_type = message.get_content_type()
    cid = _content_id(message)
    filename = message.get_filename() or (f"inline-{cid}" if cid else "attachment")

    if content_type == "text/plain" and not _is_attachment_part(message):
        text = _part_text(message).strip()
        if text and not collector.body_text:
            collector.body_text = text
        return

    if content_type == "text/html" and not _is_attachment_part(message):
        html = _part_text(message).strip()
        if html:
            collector.body_html = html
        return

    payload = _part_payload(message)
    if not payload:
        return

    extracted = ExtractedMimePart(
        filename=filename,
        content_type=content_type,
        payload=payload,
        content_id=cid,
        disposition="inline" if cid or message.get_content_disposition() == "inline" else "attachment",
    )

    if cid or extracted.disposition == "inline":
        collector.inline_parts.append(extracted)
    elif _is_attachment_part(message):
        collector.file_attachments.append(extracted)


def extract_mirror_content(raw: bytes) -> MirroredMimeContent:
    message = BytesParser(policy=policy.default).parsebytes(raw)
    collector = MirroredMimeContent()
    _walk_parts(message, collector)
    return collector