"""Mail client for SYSU Coremail (IMAP/SMTP)."""

import base64
import imaplib
import json
import os
import re
import smtplib
from datetime import datetime
from email import policy
from email.header import Header, decode_header
from email.message import EmailMessage
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.parser import BytesParser
from email.utils import parseaddr
from pathlib import Path

import click

from bin.utils import handle_errors

_HOME = Path(os.environ.get("MYTOOLKIT_HOME", Path.home() / ".mytoolkit"))
CONFIG_PATH = _HOME / "mail.json"


def load_config():
    """Load mail config."""
    if not CONFIG_PATH.exists():
        click.echo("Mail not configured. Run: mytoolkit mail config", err=True)
        raise click.Abort()
    with open(CONFIG_PATH, encoding="utf-8") as f:
        return json.load(f)


def decode_password(pwd: str) -> str:
    """Decode base64-encoded password."""
    try:
        return base64.b64decode(pwd).decode("utf-8")
    except Exception:
        return pwd


def get_imap(cfg: dict):
    """Connect to IMAP server."""
    mail = imaplib.IMAP4_SSL(cfg["imap_host"], cfg.get("imap_port", 993))
    pwd = decode_password(cfg["password"])
    mail.login(cfg["email"], pwd)
    return mail


def get_smtp(cfg: dict):
    """Connect to SMTP server."""
    smtp = smtplib.SMTP_SSL(cfg["smtp_host"], cfg.get("smtp_port", 465))
    pwd = decode_password(cfg["password"])
    smtp.login(cfg["email"], pwd)
    return smtp


def decode_mime_str(s: str) -> str:
    """Decode MIME-encoded header string (e.g. =?utf-8?...?=) to plain text."""
    if not s:
        return ""
    parts = decode_header(s)
    out = []
    for value, charset in parts:
        if isinstance(value, bytes):
            try:
                out.append(value.decode(charset or "utf-8", errors="replace"))
            except (LookupError, TypeError):
                out.append(value.decode("utf-8", errors="replace"))
        else:
            out.append(value)
    return "".join(out)


def uid_fetch(imap, uid, parts):
    """Fetch by real IMAP UID. Returns raw bytes or empty bytes if not found."""
    typ, msg_data = imap.uid("FETCH", str(uid), parts)
    if typ != "OK" or not msg_data or not msg_data[0]:
        return b""
    return msg_data[0][1]


def safe_filename(name: str) -> str:
    """Strip path separators from a filename to prevent traversal."""
    name = decode_mime_str(name).strip()
    name = name.replace("/", "_").replace("\\", "_")
    if name in ("", ".", ".."):
        return "attachment"
    return name


def _encode_filename(filename: str) -> str:
    """Return a Content-Disposition filename value compatible with Coremail.

    Coremail Web client does not reliably display RFC 2231 filename* segments.
    We use RFC 2047 encoded-word for the plain filename parameter, which older
    Chinese university webmail clients generally handle better.
    """
    ascii_name = filename.encode("ascii", "ignore").decode()
    if ascii_name == filename:
        return f'filename="{filename}"'
    encoded = Header(filename, charset="utf-8").encode()
    return f'filename="{encoded}"'


def _quote_body(text: str) -> str:
    """Prefix each line of a body with '> ' for reply quoting."""
    return "\n".join(f"> {line}" for line in text.splitlines())


def _build_reply_body(original, body: str) -> str:
    """Build a reply body that quotes the original message."""
    orig_body = ""
    b = original.get_body(preferencelist=("plain",))
    if b:
        orig_body = b.get_content() or ""
    quoted = _quote_body(orig_body) if orig_body else ""
    return (
        f"{body}\n\n"
        f"--- Original message ---\n"
        f"From: {original['From']}\n"
        f"Date: {original['Date']}\n"
        f"Subject: {original['Subject']}\n\n"
        f"{quoted}"
    )


def _references(original) -> str:
    """Build References header for a reply."""
    refs = []
    if original["References"]:
        refs.append(str(original["References"]))
    if original["Message-ID"]:
        refs.append(str(original["Message-ID"]))
    return " ".join(refs)


def _build_address_header(raw_address: str) -> str:
    """Encode a 'Display Name <addr>' header for Coremail compatibility.

    Python's compat32 policy folds long encoded words incorrectly when the
    address contains both non-ASCII display name and angle brackets. We encode
    the display name manually and preserve the raw address outside the encoded
    word.
    """
    if not raw_address:
        return ""
    display_name, addr_spec = parseaddr(raw_address)
    if not display_name or display_name == addr_spec:
        return addr_spec
    encoded = base64.b64encode(display_name.encode("utf-8")).decode("ascii")
    return f"=?utf-8?b?{encoded}?= <{addr_spec}>"


def _build_message(
    cfg: dict,
    headers: dict,
    body: str,
    html: bool,
    attach_paths,
) -> bytes:
    """Build a MIME message and return its serialized bytes.

    Uses MIMEMultipart whenever attachments are present so that non-ASCII
    display names in address headers and non-ASCII attachment filenames are
    both serialized in a Coremail-compatible way.
    """
    attach_paths = [p for p in attach_paths if p]

    if attach_paths:
        msg = MIMEMultipart("mixed")
        msg["From"] = cfg["email"]
        for key, value in headers.items():
            if key in ("To", "Cc") and value:
                msg[key] = _build_address_header(value)
            elif key == "Subject" and value:
                msg[key] = Header(value, charset="utf-8").encode()
            elif value:
                msg[key] = value

        if html:
            text_part = MIMEText("Draft includes HTML version.", "plain", "utf-8")
            html_part = MIMEText(body, "html", "utf-8")
            msg.attach(text_part)
            msg.attach(html_part)
        else:
            text_part = MIMEText(body, "plain", "utf-8")
            msg.attach(text_part)

        for path in attach_paths:
            p = Path(path)
            if not p.exists():
                click.echo(f"Attachment not found: {path}", err=True)
                continue
            with open(p, "rb") as f:
                data = f.read()
            part = MIMEApplication(data)
            part["Content-Disposition"] = f"attachment; {_encode_filename(p.name)}"
            msg.attach(part)

        return msg.as_bytes()

    # No attachments: use EmailMessage with default policy.
    msg = EmailMessage()
    msg["From"] = cfg["email"]
    for key, value in headers.items():
        if key == "Subject" and value:
            msg[key] = value
        elif value:
            msg[key] = value

    if html:
        msg.set_content("Draft includes HTML version.", subtype="plain")
        msg.add_alternative(body, subtype="html")
    else:
        msg.set_content(body)

    return msg.as_bytes()


@click.group()
def mail():
    """Mail client (IMAP/SMTP). No send command — use draft instead."""
    pass


@mail.command()
@click.option("--imap-host", default="mail.sysu.edu.cn", help="IMAP server host")
@click.option("--imap-port", default=993, type=int, help="IMAP server port")
@click.option("--smtp-host", default="mail.sysu.edu.cn", help="SMTP server host")
@click.option("--smtp-port", default=465, type=int, help="SMTP server port")
@click.option("--email", prompt="Email address", help="Email address")
@click.option("--password", prompt="Password", hide_input=True, help="Email password")
@handle_errors
def config(imap_host, imap_port, smtp_host, smtp_port, email, password):
    """Configure mail account."""
    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    cfg = {
        "imap_host": imap_host,
        "imap_port": imap_port,
        "smtp_host": smtp_host,
        "smtp_port": smtp_port,
        "email": email,
        "password": base64.b64encode(password.encode()).decode(),
    }
    with open(CONFIG_PATH, "w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=2)
    click.echo(f"Config saved to {CONFIG_PATH}")


@mail.command()
@click.option("--limit", "-n", default=20, type=int, help="Number of messages to show")
@click.option("--unseen", "-u", is_flag=True, help="Show only unseen messages")
@handle_errors
def inbox(limit, unseen):
    """List inbox messages. UID column is the real IMAP UID (stable across sessions)."""
    cfg = load_config()
    imap = get_imap(cfg)
    imap.select("INBOX", readonly=True)

    criterion = "UNSEEN" if unseen else "ALL"
    _, data = imap.uid("SEARCH", None, criterion)
    uids = data[0].split()
    uids = uids[-limit:] if len(uids) > limit else uids

    click.echo(f"{'UID':>10}  {'Date':<20}  {'From':<30}  Subject")
    click.echo("-" * 100)

    for uid in reversed(uids):
        raw_header = uid_fetch(imap, uid.decode(), "(RFC822.HEADER)")
        header = BytesParser(policy=policy.default).parsebytes(raw_header)

        subject = header["Subject"] or "(no subject)"
        from_addr = header["From"] or ""
        from_addr = re.sub(r"<.*>", "", from_addr).strip() or from_addr
        from_addr = from_addr[:28]
        date_str = ""
        if header["Date"]:
            try:
                dt = datetime.strptime(header["Date"][:31], "%a, %d %b %Y %H:%M:%S")
                date_str = dt.strftime("%Y-%m-%d %H:%M")
            except ValueError:
                date_str = str(header["Date"])[:19]

        click.echo(f"{uid.decode():>10}  {date_str:<20}  {from_addr:<30}  {subject[:50]}")

    imap.close()
    imap.logout()


@mail.command()
@click.argument("uid", type=int)
@click.option("--raw", is_flag=True, help="Show raw headers")
@handle_errors
def read(uid, raw):
    """Read a message by IMAP UID."""
    cfg = load_config()
    imap = get_imap(cfg)
    imap.select("INBOX", readonly=True)

    raw_msg = uid_fetch(imap, uid, "(RFC822)")
    msg = BytesParser(policy=policy.default).parsebytes(raw_msg)

    if raw:
        click.echo(msg.as_string())
    else:
        click.echo(f"From: {msg['From']}")
        click.echo(f"To: {msg['To']}")
        click.echo(f"Date: {msg['Date']}")
        click.echo(f"Subject: {msg['Subject']}")

        attachments = [
            safe_filename(p.get_filename() or "")
            for p in msg.walk()
            if p.get_content_disposition() == "attachment" and p.get_filename()
        ]
        if attachments:
            click.echo(f"Attachments: {', '.join(attachments)}")

        click.echo("-" * 60)
        body = msg.get_body(preferencelist=("plain",))
        if body:
            click.echo(body.get_content())
        else:
            click.echo("(no plain text body)")

    imap.close()
    imap.logout()


@mail.command()
@click.argument("uid", type=int)
@click.option("--out", "-o", "out_dir", default=".", help="Output directory (default: current dir)")
@click.option("--inline", is_flag=True, help="Also save inline parts (e.g. embedded images)")
@handle_errors
def download(uid, out_dir, inline):
    """Download attachments from a message by IMAP UID."""
    cfg = load_config()
    imap = get_imap(cfg)
    imap.select("INBOX", readonly=True)

    raw_msg = uid_fetch(imap, uid, "(RFC822)")
    imap.close()
    imap.logout()

    if not raw_msg:
        click.echo(f"Message UID {uid} not found.", err=True)
        raise click.Abort()

    msg = BytesParser(policy=policy.default).parsebytes(raw_msg)
    out = Path(out_dir).expanduser()
    out.mkdir(parents=True, exist_ok=True)

    wanted = {"attachment"}
    if inline:
        wanted.add("inline")

    saved = []
    for part in msg.walk():
        cd = part.get_content_disposition()
        if cd not in wanted:
            continue
        fn = part.get_filename()
        if not fn:
            continue
        fn = safe_filename(fn)
        target = out / fn
        # avoid overwriting existing files
        i = 1
        while target.exists():
            stem, suffix = target.stem, target.suffix
            target = out / f"{stem}-{i}{suffix}"
            i += 1
        payload = part.get_payload(decode=True)
        if not isinstance(payload, (bytes, bytearray)):
            continue
        target.write_bytes(payload)
        saved.append((target, len(payload), cd))

    if not saved:
        click.echo(f"No attachments in UID {uid}.")
        return

    for target, size, cd in saved:
        marker = "[inline]" if cd == "inline" else "        "
        click.echo(f"{marker} {target}  ({size} bytes)")


@mail.command()
@click.option("--to", required=True, help="Recipient address")
@click.option("--cc", help="CC addresses (comma-separated)")
@click.option("--subject", "-s", required=True, help="Subject")
@click.option("--body", "-b", required=True, help="Body text (\\n expanded)")
@click.option("--attach", multiple=True, help="Attachment path(s)")
@click.option("--html", is_flag=True, help="Body is HTML")
@handle_errors
def draft(to, cc, subject, body, attach, html):
    """Create a draft email (does NOT send). Saves to drafts folder."""
    cfg = load_config()
    body = body.replace("\\n", "\n").replace("\\t", "\t")

    headers = {"To": to, "Subject": subject}
    if cc:
        headers["Cc"] = cc

    raw = _build_message(cfg, headers, body, html, attach)

    imap = get_imap(cfg)
    imap.select("INBOX")
    imap.append("Drafts", None, None, raw)
    imap.close()
    imap.logout()

    click.echo("Draft saved to Drafts folder.")
    click.echo(f"  To: {to}")
    click.echo(f"  Subject: {subject}")


@mail.command()
@click.argument("uid", type=int)
@click.option("--body", "-b", required=True, help="Reply body")
@click.option("--attach", multiple=True, help="Attachment path(s)")
@click.option("--html", is_flag=True, help="Body is HTML")
@handle_errors
def reply(uid, body, attach, html):
    """Reply to a message by IMAP UID (draft mode — saves to Drafts)."""
    cfg = load_config()
    body = body.replace("\\n", "\n").replace("\\t", "\t")

    imap = get_imap(cfg)
    imap.select("INBOX", readonly=True)

    raw_orig = uid_fetch(imap, uid, "(RFC822)")
    if not raw_orig:
        imap.close()
        imap.logout()
        click.echo(f"Message UID {uid} not found.", err=True)
        raise click.Abort()
    original = BytesParser(policy=policy.default).parsebytes(raw_orig)

    subject = original["Subject"] or ""
    if not str(subject).startswith("Re:"):
        subject = f"Re: {subject}"

    headers = {
        "To": original["Reply-To"] or original["From"],
        "Subject": subject,
        "In-Reply-To": original["Message-ID"],
        "References": _references(original),
    }

    content = _build_reply_body(original, body)
    if html:
        # For HTML replies we still include the plain quoted body as plain text
        # alternative and an HTML rendering as primary alternative.
        content = _build_reply_body(original, body)
        raw = _build_message(cfg, headers, content.replace("\n", "<br>\n"), True, attach)
    else:
        raw = _build_message(cfg, headers, content, False, attach)

    imap.append("Drafts", None, None, raw)
    imap.close()
    imap.logout()

    click.echo(f"Reply draft saved for UID {uid}.")


@mail.command()
@click.argument("uid", type=int)
@click.option("--to", required=True, help="Forward to")
@click.option("--body", "-b", default="", help="Additional body text")
@click.option("--attach", multiple=True, help="Attachment path(s)")
@click.option("--html", is_flag=True, help="Body is HTML")
@handle_errors
def forward(uid, to, body, attach, html):
    """Forward a message by IMAP UID (draft mode — saves to Drafts)."""
    cfg = load_config()
    body = body.replace("\\n", "\n").replace("\\t", "\t")

    imap = get_imap(cfg)
    imap.select("INBOX", readonly=True)

    raw_orig = uid_fetch(imap, uid, "(RFC822)")
    if not raw_orig:
        imap.close()
        imap.logout()
        click.echo(f"Message UID {uid} not found.", err=True)
        raise click.Abort()
    original = BytesParser(policy=policy.default).parsebytes(raw_orig)

    orig_body = ""
    b = original.get_body(preferencelist=("plain",))
    if b:
        orig_body = b.get_content()

    subject = f"Fwd: {original['Subject']}"
    content = (
        f"{body}\n\n"
        f"--- Forwarded message ---\n"
        f"From: {original['From']}\n"
        f"Date: {original['Date']}\n"
        f"Subject: {original['Subject']}\n\n"
        f"{orig_body}"
    )

    headers = {"To": to, "Subject": subject}
    if html:
        raw = _build_message(cfg, headers, content.replace("\n", "<br>\n"), True, attach)
    else:
        raw = _build_message(cfg, headers, content, False, attach)

    imap.append("Drafts", None, None, raw)
    imap.close()
    imap.logout()

    click.echo(f"Forward draft saved for UID {uid}.")