from __future__ import annotations

import asyncio
import base64
import email
import imaplib
import logging
import re
import socket
import time
import urllib.parse
from datetime import datetime, timezone

from mail.oauth import OAuthError, access_token, drop_cached_token, xoauth2_string

logger = logging.getLogger(__name__)

_IMAP_MONTHS = (
    "Jan", "Feb", "Mar", "Apr", "May", "Jun",
    "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
)


def imap_search_date(when: datetime) -> str:
    return f"{when.day:02d}-{_IMAP_MONTHS[when.month - 1]}-{when.year}"

_CODE_NEAR_KEYWORD = re.compile(
    r"\b(?:security\s*code|verification\s*code|verify(?:\s*code)?|code)\D{0,20}?(\d{4,8})\b",
    re.IGNORECASE,
)
_ANY_CODE = re.compile(r"\b(\d{4,8})\b")


def _get_text_body(msg: email.message.Message) -> str:
    if msg.is_multipart():
        parts = []
        for part in msg.walk():
            if part.get_content_type() == "text/plain" and not part.get_filename():
                try:
                    payload = part.get_payload(decode=True) or b""
                    charset = part.get_content_charset() or "utf-8"
                    parts.append(payload.decode(charset, errors="replace"))
                except Exception:
                    continue
        return "\n".join(parts)

    try:
        payload = msg.get_payload(decode=True)
        if payload is None:
            return str(msg.get_payload())
        charset = msg.get_content_charset() or "utf-8"
        return payload.decode(charset, errors="replace")
    except Exception:
        return str(msg.get_payload())


def extract_code(raw_email_bytes: bytes) -> str | None:
    msg = email.message_from_bytes(raw_email_bytes)
    subject = msg.get("Subject", "") or ""
    body = _get_text_body(msg)
    text = f"{subject}\n{body}"

    match = _CODE_NEAR_KEYWORD.search(text)
    if match:
        return match.group(1)

    match = _ANY_CODE.search(text)
    if match:
        return match.group(1)

    return None


def imap_cfg_from_instance(cfg) -> dict | None:
    host = getattr(cfg, "imap_host", None)
    user = getattr(cfg, "imap_user", None)
    if not host or not user:
        return None
    out = {
        "host": host,
        "port": getattr(cfg, "imap_port", None) or 993,
        "username": user,
        "password": getattr(cfg, "imap_password", None),
    }
    refresh = getattr(cfg, "imap_refresh_token", None)
    if refresh:
        out["auth"] = "xoauth2"
        out["provider"] = getattr(cfg, "imap_provider", None)
        out["refresh_token"] = refresh
    proxy = getattr(cfg, "proxy_url", None)
    if proxy:
        out["proxy"] = proxy
    return out


def _uses_xoauth2(imap_cfg: dict) -> bool:
    return imap_cfg.get("auth") == "xoauth2" or bool(imap_cfg.get("refresh_token"))


def _authenticate(conn: imaplib.IMAP4, imap_cfg: dict) -> None:
    user = imap_cfg["username"]
    if _uses_xoauth2(imap_cfg):
        token = access_token(imap_cfg)
        sasl = xoauth2_string(user, token)

        def _auth(_challenge):
            return sasl

        try:
            conn.authenticate("XOAUTH2", _auth)
        except imaplib.IMAP4.error:
            drop_cached_token(imap_cfg.get("refresh_token") or "")
            raise
        return
    conn.login(user, imap_cfg["password"])


def _http_connect_tunnel(proxy_url: str, host: str, port: int, timeout) -> socket.socket:
    parsed = urllib.parse.urlparse(proxy_url)
    if parsed.scheme != "http":
        raise OSError(f"unsupported IMAP proxy scheme: {parsed.scheme!r}")
    if not parsed.hostname:
        raise OSError("proxy URL is missing a host")
    proxy_port = parsed.port or 80
    sock = socket.create_connection((parsed.hostname, proxy_port), timeout)
    try:
        target = f"{host}:{int(port)}"
        lines = [f"CONNECT {target} HTTP/1.1", f"Host: {target}"]
        if parsed.username is not None:
            user = urllib.parse.unquote(parsed.username)
            password = urllib.parse.unquote(parsed.password or "")
            token = base64.b64encode(f"{user}:{password}".encode()).decode("ascii")
            lines.append(f"Proxy-Authorization: Basic {token}")
        lines.extend(["", ""])
        sock.sendall("\r\n".join(lines).encode("ascii"))
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = sock.recv(4096)
            if not chunk:
                raise OSError("proxy closed during CONNECT")
            buf += chunk
            if len(buf) > 65536:
                raise OSError("oversized CONNECT response")
        status = buf.split(b"\r\n", 1)[0].decode("ascii", "replace")
        parts = status.split()
        if len(parts) < 2 or not parts[1].isdigit() or not parts[1].startswith("2"):
            raise OSError(f"proxy CONNECT failed: {status}")
        return sock
    except BaseException:
        sock.close()
        raise


class _ProxiedIMAP4_SSL(imaplib.IMAP4_SSL):
    def __init__(self, host, port, timeout, proxy_url):
        self._proxy_url = proxy_url
        super().__init__(host, port, timeout=timeout)

    def _create_socket(self, timeout):
        sock = _http_connect_tunnel(self._proxy_url, self.host, self.port, timeout)
        return self.ssl_context.wrap_socket(sock, server_hostname=self.host)


def _open_connection(imap_cfg: dict) -> imaplib.IMAP4_SSL:
    host = imap_cfg["host"]
    port = imap_cfg.get("port") or 993
    connect_timeout = imap_cfg.get("timeout", 10)
    proxy = imap_cfg.get("proxy")
    if proxy:
        conn = _ProxiedIMAP4_SSL(host, port, connect_timeout, proxy)
    else:
        conn = imaplib.IMAP4_SSL(host, port, timeout=connect_timeout)
    try:
        _authenticate(conn, imap_cfg)
        return conn
    except Exception:
        try:
            conn.logout()
        except Exception:
            pass
        raise


def validate_imap_cfg(imap_cfg: dict) -> bool:
    try:
        conn = _open_connection(imap_cfg)
    except (OSError, imaplib.IMAP4.error, RuntimeError, OAuthError):
        return False
    try:
        conn.logout()
    except OSError:
        pass
    return True


async def wait_for_code(
    imap_cfg: dict,
    since_seen: set[bytes],
    timeout: float = 30,
    poll: float = 3,
) -> str | None:
    since_date = imap_search_date(datetime.now(timezone.utc))
    start = time.monotonic()

    while time.monotonic() - start < timeout:
        try:
            code = await asyncio.to_thread(_poll_once, imap_cfg, since_seen, since_date)
        except Exception:
            logger.warning("imap poll cycle failed, will retry", exc_info=True)
            code = None
        if code:
            return code
        await asyncio.sleep(poll)

    return None


def _poll_once(imap_cfg: dict, since_seen: set[bytes], since_date: str) -> str | None:
    mailbox = imap_cfg.get("mailbox", "INBOX")
    conn = _open_connection(imap_cfg)
    try:
        conn.select(mailbox)

        status, data = conn.search(
            None,
            f'(UNSEEN SINCE "{since_date}" FROM "microsoft.com")',
        )
        if status != "OK" or not data or not data[0]:
            return None

        for msg_id in data[0].split():
            if msg_id in since_seen:
                continue

            status, msg_data = conn.fetch(msg_id, "(RFC822)")
            if status != "OK" or not msg_data or not msg_data[0]:
                continue

            raw = msg_data[0][1]
            since_seen.add(msg_id)

            code = extract_code(raw)
            if code:
                return code

        return None
    finally:
        try:
            conn.close()
        except Exception:
            pass
        try:
            conn.logout()
        except Exception:
            pass