"""IMAP IDLE stream handler for real-time email notifications."""

import asyncio
import email
import imaplib
import logging
import time
from collections.abc import AsyncIterator
from typing import Any

from flowire_sdk import StreamHandler

from fw_nodes_email.nodes.read_emails import _decode_header_value, _extract_attachments, _extract_body

logger = logging.getLogger(__name__)

# IMAP IDLE timeout — RFC 2177 recommends re-issuing IDLE every 29 minutes.
# We use 25 minutes to stay safely within server limits.
IDLE_TIMEOUT_SECONDS = 25 * 60


def _parse_email_message(raw_bytes: bytes) -> dict[str, Any]:
    """Parse a raw email into a structured dict."""
    from email.utils import parsedate_to_datetime

    msg = email.message_from_bytes(raw_bytes)
    plain_body, html_body = _extract_body(msg)
    attachments = _extract_attachments(msg, include_content=False)

    date_str = msg.get("Date")
    date_iso = ""
    if date_str:
        try:
            date_iso = parsedate_to_datetime(date_str).isoformat()
        except Exception:
            date_iso = date_str

    return {
        "message_id": msg.get("Message-ID", ""),
        "from": _decode_header_value(msg.get("From")),
        "to": _decode_header_value(msg.get("To")),
        "cc": _decode_header_value(msg.get("Cc")),
        "subject": _decode_header_value(msg.get("Subject")),
        "date": date_iso,
        "body_plain": plain_body,
        "body_html": html_body,
        "attachments": attachments,
    }


class ImapIdleStreamHandler(StreamHandler):
    """Handler for IMAP IDLE push notifications.

    Connects to an IMAP server and uses the IDLE command to receive
    real-time notifications when new emails arrive.

    Stream type: imap_idle
    """

    STREAM_TYPE = "imap_idle"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._imap: imaplib.IMAP4_SSL | imaplib.IMAP4 | None = None
        self._stop_idle = False

    async def connect(self) -> None:
        """Establish IMAP connection and select folder."""
        credential = await self.resolve_credential()

        host = credential["host"]
        port = credential.get("port", 993)
        use_ssl = credential.get("use_ssl", True)

        readonly = credential.get("readonly", False)

        config = self.subscription.connection_config or {}
        folder = config.get("folder", "INBOX")

        logger.info(
            "Connecting to IMAP: %s:%s folder=%s subscription=%s",
            host,
            port,
            folder,
            self.subscription.id,
        )

        def _connect_sync():
            if use_ssl:
                imap = imaplib.IMAP4_SSL(host, port)
            else:
                imap = imaplib.IMAP4(host, port)

            imap.login(credential["username"], credential["password"])
            imap.select(folder, readonly=readonly)

            # Check IDLE capability
            caps = imap.capability()
            if caps[0] == "OK":
                cap_str = b" ".join(caps[1]).upper()
                if b"IDLE" not in cap_str:
                    raise RuntimeError(f"IMAP server {host} does not support IDLE")

            return imap

        self._imap = await asyncio.to_thread(_connect_sync)

        logger.info(
            "Connected to IMAP: %s subscription=%s",
            host,
            self.subscription.id,
        )

    async def disconnect(self) -> None:
        """Close IMAP connection.

        Shuts down the raw socket first to unblock any thread that may be
        blocking on imap.readline() (e.g., _idle_wait during IDLE mode).
        Without this, imap.close()/logout() can deadlock because imaplib
        is not thread-safe and the _idle_wait thread holds the socket.
        """
        self._stop_idle = True
        if self._imap:
            # Shut down the raw socket first to unblock _idle_wait thread
            try:
                self._imap.sock.shutdown(2)  # SHUT_RDWR
            except Exception:
                pass
            try:
                await asyncio.to_thread(self._imap.close)
            except Exception:
                pass
            try:
                await asyncio.to_thread(self._imap.logout)
            except Exception:
                pass
            self._imap = None

    async def listen(self) -> AsyncIterator[dict[str, Any]]:
        """Yield events when new emails arrive via IMAP IDLE.

        Uses IDLE to wait for EXISTS responses (new messages).
        Periodically re-issues IDLE per RFC 2177.
        """
        if not self._imap:
            raise RuntimeError("Not connected")

        self._stop_idle = False

        # Track the highest UID we've seen so we only yield genuinely new messages
        last_uid = await asyncio.to_thread(self._get_latest_uid)

        while not self._stop_idle:
            try:
                # Enter IDLE mode and wait for server notifications
                new_mail = await asyncio.to_thread(self._idle_wait)

                if self._stop_idle:
                    break

                if new_mail:
                    # Fetch new messages since last known UID
                    new_emails, last_uid = await asyncio.to_thread(self._fetch_new_since_uid, last_uid)

                    for email_data in new_emails:
                        yield {
                            "event": "new_email",
                            "data": email_data,
                        }

            except (imaplib.IMAP4.error, imaplib.IMAP4.abort, OSError) as e:
                logger.error("IMAP IDLE error: %s subscription=%s", e, self.subscription.id)
                raise

    def _get_latest_uid(self) -> int:
        """Get the highest UID in the current folder."""
        imap = self._imap
        status, data = imap.uid("search", None, "ALL")
        if status != "OK" or not data[0]:
            return 0
        uids = data[0].split()
        return int(uids[-1]) if uids else 0

    def _fetch_new_since_uid(self, last_uid: int) -> tuple[list[dict[str, Any]], int]:
        """Fetch emails with UID greater than last_uid. Returns (emails, new_last_uid)."""
        imap = self._imap
        # Search for UIDs greater than our last seen
        search_criteria = f"UID {last_uid + 1}:*"
        status, data = imap.uid("search", None, search_criteria)

        if status != "OK" or not data[0]:
            return [], last_uid

        uids = data[0].split()
        # Filter out the last_uid itself (IMAP range is inclusive)
        uids = [u for u in uids if int(u) > last_uid]

        if not uids:
            return [], last_uid

        emails = []
        new_last_uid = last_uid

        for uid in uids:
            status, msg_data = imap.uid("fetch", uid, "(RFC822)")
            if status != "OK" or not msg_data[0]:
                continue

            raw_email = msg_data[0][1]
            email_data = _parse_email_message(raw_email)
            email_data["uid"] = int(uid)
            emails.append(email_data)

            uid_int = int(uid)
            if uid_int > new_last_uid:
                new_last_uid = uid_int

        return emails, new_last_uid

    def _idle_wait(self) -> bool:
        """Send IDLE command, wait for new mail, then return.

        Returns True if new mail was indicated, False on timeout.
        """
        imap = self._imap

        # Send IDLE command
        # imaplib has no public IDLE API — _new_tag() generates the required
        # command tag (e.g. b"ABCD1"). This is the standard workaround; see
        # CPython's imaplib source. Monitor for breakage on Python upgrades.
        tag = imap._new_tag()
        imap.send(tag + b" IDLE\r\n")

        # Wait for continuation response (+)
        response = imap.readline()
        if not response.startswith(b"+"):
            # Server didn't accept IDLE, fall back to polling
            logger.warning("IMAP server did not accept IDLE: %s", response)
            imap.send(b"DONE\r\n")
            imap.readline()
            time.sleep(30)
            return True  # Check anyway

        # Wait for untagged responses (EXISTS = new mail)
        new_mail = False
        imap.sock.settimeout(IDLE_TIMEOUT_SECONDS)
        try:
            while not self._stop_idle:
                line = imap.readline()
                if not line:
                    break
                line_str = line.decode("utf-8", errors="replace").strip()
                if "EXISTS" in line_str:
                    new_mail = True
                    break
                if "EXPUNGE" in line_str:
                    # Message deleted, not new mail but worth re-checking
                    break
        except TimeoutError:
            # Normal IDLE timeout, just re-issue
            pass
        except OSError:
            # Socket timeout
            pass
        finally:
            imap.sock.settimeout(None)

        # End IDLE
        imap.send(b"DONE\r\n")
        # Read the tagged response to IDLE
        try:
            imap.readline()
        except Exception:
            pass

        return new_mail
