from __future__ import annotations

import email
import mailbox
import typing
from email.utils import parsedate_to_datetime

from ..base import BaseCodec, BaseFileIterable


class MBOXIterable(BaseFileIterable):
    def __init__(
        self,
        filename: str = None,
        stream: typing.IO = None,
        codec: BaseCodec = None,
        mode: str = "r",
        encoding: str = "utf8",
        options: dict = None,
    ):
        if options is None:
            options = {}
        super().__init__(filename, stream, codec=codec, binary=False, mode=mode, encoding=encoding, options=options)
        self.reset()
        pass

    def reset(self):
        """Reset iterable"""
        super().reset()
        self.pos = 0
        if self.mode == "r":
            # MBOX format requires file path, not stream
            if self.filename is None:
                raise ValueError("MBOX format requires filename, not stream")

            self.mbox = mailbox.mbox(self.filename)
            self.iterator = iter(self.mbox)
        else:
            if self.filename is None:
                raise ValueError("MBOX format requires filename for writing")
            self.mbox = mailbox.mbox(self.filename, create=True)

    @staticmethod
    def id() -> str:
        return "mbox"

    @staticmethod
    def is_flatonly() -> bool:
        return False

    def _email_to_dict(self, msg: email.message.Message) -> dict:
        """Convert email message to dictionary"""
        result = {}

        # Basic headers
        for header in [
            "From",
            "To",
            "Cc",
            "Bcc",
            "Subject",
            "Date",
            "Message-ID",
            "In-Reply-To",
            "References",
            "Reply-To",
            "Sender",
        ]:
            value = msg.get(header)
            if value:
                result[header.lower()] = value

        # Parse date
        if "date" in result:
            try:
                result["date_parsed"] = parsedate_to_datetime(result["date"]).isoformat()
            except Exception:
                pass

        # Body content
        if msg.is_multipart():
            parts = []
            for part in msg.walk():
                if part.get_content_maintype() == "text":
                    parts.append(
                        {
                            "content_type": part.get_content_type(),
                            "content": part.get_payload(decode=True).decode("utf-8", errors="ignore"),
                        }
                    )
            result["body"] = parts
            result["body_text"] = "\n\n".join([p["content"] for p in parts])
        else:
            payload = msg.get_payload(decode=True)
            if payload:
                try:
                    result["body"] = payload.decode("utf-8", errors="ignore")
                except Exception:
                    result["body"] = str(payload)
            result["body_text"] = result.get("body", "")

        # All headers
        result["headers"] = dict(msg.items())

        # Attachments info
        attachments = []
        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_disposition() == "attachment":
                    attachments.append(
                        {
                            "filename": part.get_filename(),
                            "content_type": part.get_content_type(),
                            "size": len(part.get_payload(decode=True) or b""),
                        }
                    )
        if attachments:
            result["attachments"] = attachments

        return result

    def read(self) -> dict:
        """Read single MBOX record"""
        msg = next(self.iterator)
        self.pos += 1
        return self._email_to_dict(msg)

    def read_bulk(self, num: int = 10) -> list[dict]:
        """Read bulk MBOX records"""
        chunk = []
        for _n in range(0, num):
            try:
                chunk.append(self.read())
            except StopIteration:
                break
        return chunk

    def write(self, record: dict):
        """Write single MBOX record"""
        msg = email.message.EmailMessage()

        # Set headers
        for key, value in record.items():
            if key.lower() in [
                "from",
                "to",
                "cc",
                "bcc",
                "subject",
                "date",
                "message-id",
                "in-reply-to",
                "references",
                "reply-to",
                "sender",
            ]:
                msg[key] = str(value)
            elif key == "body" or key == "body_text":
                msg.set_content(str(value))

        self.mbox.add(msg)
        self.mbox.flush()

    def write_bulk(self, records: list[dict]):
        """Write bulk MBOX records"""
        for record in records:
            self.write(record)

    def close(self):
        """Close MBOX file"""
        if hasattr(self, "mbox"):
            self.mbox.close()
        super().close()
