#!/usr/bin/env python3
"""mail/import - pull OnlyOffice mails into var/mail/ as markdown.

  bin/mail/import --from-raw var/mail   convert Go-synced message.json to md
  bin/mail/import                       import newest inbox messages
  bin/mail/import --folder sent         import sent folder
  bin/mail/import --since 2026-01-01    only messages after a date
  bin/mail/import --limit 50            cap messages per run
  bin/mail/import --no-attachments      body only, skip attachment conversion
  bin/mail/import --ocr                OCR scanned PDFs/images via docling
  bin/mail/import --dry-run            list messages without writing anything

Writes one directory per message: var/mail/{folder}/{message_id}/
  message.md                frontmatter + markdown body
  attachments/              raw attachment files (zips unpacked to _unpacked/)
  attachments/*.md          converted attachment content

Indexing is a separate step (`bin/brain/index.go --rebuild`): conversion can
crash in native docling and must not leave the brain DB mid-transaction.

Requires ONLYOFFICE_URL/USER/PASS in .env (or env). Idempotent: a message
already present (message.md exists) is skipped unless --force.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
import urllib.parse
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "bin" / "tools"))

from mailconv import (  # noqa: E402
    ARCHIVE_SUFFIXES,
    IMAGE_SUFFIXES,
    LEGACY_OFFICE_SUFFIXES,
    TEXT_SUFFIXES,
    html_to_markdown,
    is_convertible,
    normalize_markdown,
    subject_to_filename,
    zip_extract_safe,
)

import requests  # noqa: E402

FOLDER_IDS = {"inbox": 1, "sent": 2, "drafts": 3, "trash": 4, "spam": 5}
DEFAULT_LIMIT = 25


def load_env() -> dict:
    env = {k: v for k, v in os.environ.items()}
    envfile = ROOT / ".env"
    if envfile.exists():
        for line in envfile.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, _, v = line.partition("=")
            env.setdefault(k.strip(), v.strip().strip("\"'"))
    url = env.get("ONLYOFFICE_URL") or env.get("OO_URL")
    user = env.get("ONLYOFFICE_USER") or env.get("OO_USER")
    password = env.get("ONLYOFFICE_PASS") or env.get("OO_PASSWORD")
    missing = [n for n, v in (("ONLYOFFICE_URL", url), ("ONLYOFFICE_USER", user),
                              ("ONLYOFFICE_PASS", password)) if not v]
    if missing:
        sys.exit(f"mail/import: missing {', '.join(missing)} (need .env or env)")
    return {"url": url.rstrip("/"), "user": user, "password": password}


class OOClient:
    def __init__(self, conf: dict):
        self.base = conf["url"]
        self.session = requests.Session()
        self.token = None
        self._login(conf)

    def _login(self, conf: dict) -> None:
        r = self.session.post(f"{self.base}/api/2.0/authentication.json",
                              json={"userName": conf["user"], "password": conf["password"], "type": 0},
                              timeout=30)
        r.raise_for_status()
        body = r.json()
        self.token = (body.get("response") or {}).get("token", "")
        if not self.token:
            sys.exit("mail/import: authentication failed (empty token)")

    def _headers(self) -> dict:
        return {"Authorization": f"Bearer {self.token}", "Accept": "application/json"}

    def get(self, path: str, params: dict | None = None):
        r = self.session.get(f"{self.base}{path}", params=params, headers=self._headers(), timeout=30)
        r.raise_for_status()
        return r.json()

    def list_messages(self, folder: int, page: int = 1, count: int = DEFAULT_LIMIT) -> list[dict]:
        data = self.get("/api/2.0/mail/messages",
                        params={"folder": folder, "page": page, "count": count})
        return data.get("response", [])

    def get_message(self, message_id: str) -> dict:
        data = self.get(f"/api/2.0/mail/messages/{message_id}")
        return data.get("response", {})

    def download_attachment(self, attach_id, dest: Path) -> bool:
        """Download one attachment via the portal session cookie (.ashx handler)."""
        url = f"{self.base}/addons/mail/httphandlers/download.ashx?attachid={attach_id}"
        r = self.session.get(url, timeout=60)
        if r.status_code != 200:
            return False
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(r.content)
        return True


def folder_id(name: str) -> int:
    if name in FOLDER_IDS:
        return FOLDER_IDS[name]
    if name.isdigit():
        return int(name)
    sys.exit(f"mail/import: unknown folder '{name}' (use {', '.join(FOLDER_IDS)})")


def safe_attachment_name(att: dict) -> str:
    name = att.get("fileName") or att.get("storedName") or "attachment"
    name = re.sub(r"[^\w.\- ]+", "_", name)
    return name


def convert_file_to_md(path: Path, ocr: bool) -> str | None:
    """Convert one attachment file to markdown text; None when not convertible."""
    suffix = path.suffix.lower()
    if suffix in TEXT_SUFFIXES:
        return normalize_markdown(path.read_text(encoding="utf-8", errors="replace"))
    if suffix in (".docx", ".pptx", ".xlsx", ".html", ".htm", ".epub", ".eml", ".msg"):
        try:
            from markitdown import MarkItDown
            md = MarkItDown()
            result = md.convert(str(path))
            return normalize_markdown(result.text_content)
        except Exception as e:
            return f"\n<!-- conversion failed: {e} -->\n"
    if suffix == ".pdf":
        return _convert_pdf(path, ocr)
    if suffix in IMAGE_SUFFIXES and ocr:
        return _convert_pdf(path, ocr)
    if suffix in LEGACY_OFFICE_SUFFIXES:
        return _convert_legacy(path)
    if suffix in ARCHIVE_SUFFIXES:
        return None  # handled by caller (unpack + recurse)
    return None


def _convert_pdf(path: Path, ocr: bool) -> str:
    """Convert one PDF to markdown.

    Fast path: poppler's pdftotext (-layout) extracts exact text from
    born-digital PDFs in ~15ms vs docling's 1-3s. Only textless PDFs (scanned
    pages, layout-heavy) fall back to docling, which runs isolated in a
    subprocess because its native onnx/RT-DETR has segfaulted the main process.
    """
    text = _pdf_fast_text(path)
    if ocr or text is None or not text.strip():
        return _convert_pdf_docling(path, ocr)
    return normalize_markdown(text)


def _pdf_fast_text(path: Path) -> str | None:
    """pdftotext -layout; None when poppler is unavailable (or the PDF has no text layer)."""
    try:
        proc = subprocess.run(
            ["pdftotext", "-layout", str(path), "-"],
            capture_output=True, timeout=60)
    except (OSError, subprocess.TimeoutExpired):
        return None
    if proc.returncode != 0:
        return None
    return proc.stdout.decode("utf-8", errors="replace")


def _convert_pdf_docling(path: Path, ocr: bool) -> str:
    try:
        proc = subprocess.run(
            [sys.executable, os.path.abspath(__file__), "--pdf-worker", str(path),
             "--ocr" if ocr else "--no-ocr"],
            capture_output=True, text=True, timeout=600)
    except subprocess.TimeoutExpired:
        return "\n<!-- pdf conversion timed out -->\n"
    if proc.returncode != 0:
        tail = proc.stderr.strip().splitlines()[-3:]
        return f"\n<!-- pdf conversion failed: {proc.returncode}: {' | '.join(tail)} -->\n"
    return proc.stdout


def _pdf_worker(path: Path, ocr: bool) -> None:
    """docling worker entry: prints converted markdown on stdout, exits non-zero on error."""
    try:
        from docling.document_converter import DocumentConverter, PdfFormatOption
        from docling.datamodel.pipeline_options import PdfPipelineOptions
        opts = PdfPipelineOptions()
        opts.do_ocr = bool(ocr)
        opts.do_table_structure = True
        conv = DocumentConverter(format_options={"pdf": PdfFormatOption(pipeline_options=opts)})
        res = conv.convert(str(path))
        sys.stdout.write(normalize_markdown(res.document.export_to_markdown()))
        sys.exit(0)
    except Exception as e:
        # errors/stacktraces to stderr; the caller only reports a one-liner
        print(f"pdf-worker: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc(file=sys.stderr)
        sys.exit(1)


def _convert_legacy(path: Path) -> str:
    """Legacy .doc/.xls/.ppt -> md via pandoc (installed) or a stub."""
    try:
        out = subprocess.run(["pandoc", str(path), "-t", "markdown"],
                             capture_output=True, text=True, timeout=120)
        if out.returncode == 0 and out.stdout.strip():
            return normalize_markdown(out.stdout)
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return f"\n<!-- legacy {path.suffix} not convertible (pandoc unavailable) -->\n"


def write_message_md(msg: dict, folder: str, out_dir: Path, target_dir: Path | None = None) -> Path:
    import yaml
    body_html = msg.get("htmlBody") or ""
    body_text = msg.get("textBody") or ""
    body_md = ""
    if body_html.strip():
        body_md = html_to_markdown(body_html)
    elif body_text.strip():
        body_md = normalize_markdown(body_text)
    # accept both OnlyOffice (receivedDate) and Go-sync (receivedAt) date keys
    date = msg.get("receivedDate") or msg.get("receivedAt") or ""
    if date and not isinstance(date, str):
        date = str(date)
    meta = {
        "id": msg.get("id"),
        "source": msg.get("source"),
        "folder": folder,
        "subject": msg.get("subject", ""),
        "from": msg.get("from", ""),
        "to": msg.get("to", ""),
        "cc": msg.get("cc", ""),
        "date": date,
        "has_attachments": bool(msg.get("hasAttachments")),
        "mime_message_id": msg.get("mimeMessageId", ""),
        "calendar_uid": msg.get("calendarUid", ""),
        "type": "mail",
    }
    meta = {k: v for k, v in meta.items() if v not in (None, "")}
    frontmatter = "---\n" + yaml.safe_dump(meta, sort_keys=False, allow_unicode=True).strip() + "\n---\n"
    content = f"{frontmatter}\n# {meta.get('subject','')}\n\n{body_md}".strip() + "\n"
    if target_dir is not None:
        msg_dir = target_dir
    else:
        msg_dir = out_dir / folder / str(meta.get("id"))
    msg_dir.mkdir(parents=True, exist_ok=True)
    md_path = msg_dir / "message.md"
    md_path.write_text(content, encoding="utf-8")
    return md_path


def convert_attachments(msg: dict, msg_dir: Path, ocr: bool) -> list[dict]:
    """Download + convert each attachment; returns [{name, md, raw}] summaries.

    raw file keeps the API storedName (unique hash, avoids collisions); the
    markdown is named after the friendly fileName when available.

    In --from-raw mode attachments are already on disk (Go sync wrote them;
    .ics already has a structured .md sidecar). Files with an existing .md
    sidecar are left as-is, only unconverted raws are converted here.
    """
    out: list[dict] = []
    atts = msg.get("attachments") or []
    att_dir = msg_dir / "attachments"
    for att in atts:
        aid = att.get("fileId")
        display = safe_attachment_name(att)
        stored = att.get("storedName")
        raw_name = safe_attachment_name({"storedName": stored}) if stored else display
        raw = att_dir / raw_name
        if aid and not raw.exists() and OOCLIENT is not None:
            if not OOCLIENT.download_attachment(aid, raw):
                out.append({"name": display, "md": "\n<!-- download failed -->\n", "raw": str(raw)})
                continue
        if not raw.exists():
            out.append({"name": display, "md": "\n<!-- raw missing -->\n", "raw": str(raw)})
            continue
        md_stem = Path(display).stem or raw.stem
        md_path = att_dir / f"{md_stem}.md"
        # Go sync pre-wrote structured .md for .ics; keep it.
        if not md_path.exists():
            md_text = _convert_att_recursive(raw, ocr)
            md_path.write_text(f"# Attachment: {display}\n\n{md_text}\n", encoding="utf-8")
        else:
            md_text = md_path.read_text(encoding="utf-8", errors="replace")
        out.append({"name": display, "md": md_text, "raw": str(raw), "md_file": str(md_path)})
    return out


def _convert_att_recursive(path: Path, ocr: bool) -> str:
    if path.suffix.lower() in ARCHIVE_SUFFIXES:
        parts: list[str] = []
        unpack = path.parent / "_unpacked" / path.stem
        files = zip_extract_safe(path, unpack)
        for f in files:
            sub = _convert_att_recursive(f, ocr)
            if sub and sub.strip():
                parts.append(f"## {f.name}\n\n{sub}")
        return "\n\n".join(parts) if parts else "\n<!-- empty zip -->\n"
    text = convert_file_to_md(path, ocr)
    return text or "\n<!-- not convertible -->\n"


# module-level client for attachment downloads in convert_attachments
OOCLIENT: OOClient | None = None


def convert_one(msg_dir: Path, full: dict, folder: str, out_root: Path,
                ocr: bool, no_attachments: bool, target_dir: Path | None = None) -> dict:
    """Write message.md + convert attachments for one message dict.

    Works for both live API messages and the Go-sync message.json shape
    (source field optional; attachments read from attachments/ dir).
    target_dir overrides the derived path (used by --from-raw where the
    directory layout is authoritative, not the message folder field).
    """
    mid = str(full.get("id"))
    write_message_md(full, folder, out_root, target_dir=target_dir)
    converted: list[dict] = []
    if not no_attachments:
        converted = convert_attachments(full, msg_dir, ocr)
    return {"id": mid, "subject": full.get("subject", ""),
            "date": full.get("receivedDate", "") or full.get("receivedAt", ""),
            "attachments": len(converted)}


def main(argv: list[str]) -> int:
    global OOCLIENT
    p = argparse.ArgumentParser(description="pull OnlyOffice mails to var/mail as markdown")
    p.add_argument("--folder", default="inbox", help="inbox|sent|drafts|trash|spam or numeric id")
    p.add_argument("--limit", type=int, default=DEFAULT_LIMIT, help="max messages per run")
    p.add_argument("--offset", type=int, default=0, help="skip N messages")
    p.add_argument("--since", default="", help="only messages received after YYYY-MM-DD")
    p.add_argument("--id", action="append", default=[], help="import specific message id (repeatable)")
    p.add_argument("--from-raw", default="",
                   help="convert Go-synced dirs (var/mail/<folder>/<id>/message.json) to markdown")
    p.add_argument("--no-attachments", action="store_true", help="skip attachment download+convert")
    p.add_argument("--ocr", action="store_true", help="OCR scanned PDFs/images via docling")
    p.add_argument("--force", action="store_true", help="re-import even if message.md exists")
    p.add_argument("--dry-run", action="store_true", help="list messages, write nothing")
    p.add_argument("--json", action="store_true")
    p.add_argument("--pdf-worker", default="", help=argparse.SUPPRESS)
    p.add_argument("--no-ocr", action="store_true", help=argparse.SUPPRESS)
    a = p.parse_args(argv)

    if a.pdf_worker:
        _pdf_worker(Path(a.pdf_worker), ocr=not a.no_ocr)
        return 0

    conf = load_env()
    fid = folder_id(a.folder)
    out_root = ROOT / "var" / "mail"
    summary: list[dict] = []
    if a.from_raw:
        OOCLIENT = None
        raw_root = Path(a.from_raw)
        for msg_dir in sorted(raw_root.rglob("message.json")):
            mid = msg_dir.parent.name
            entry = {"id": mid, "subject": "", "date": "",
                     "attachments": 0, "skipped": False}
            md_path = msg_dir.parent / "message.md"
            if md_path.exists() and not a.force:
                entry["skipped"] = True
                summary.append(entry)
                continue
            if a.dry_run:
                entry["skipped"] = "dry-run"
                summary.append(entry)
                continue
            full = json.loads(msg_dir.read_text(encoding="utf-8"))
            entry.update(convert_one(msg_dir.parent, full, full.get("folder") or a.folder,
                                     raw_root, a.ocr, a.no_attachments,
                                     target_dir=msg_dir.parent))
            summary.append(entry)
    else:
        OOCLIENT = OOClient(conf)
        if a.id:
            messages = [{"id": i} for i in a.id]
        else:
            page = 1
            messages = []
            want = a.offset + a.limit
            while len(messages) < want:
                count = min(DEFAULT_LIMIT, want - len(messages))
                chunk = OOCLIENT.list_messages(fid, page=page, count=count)
                if not chunk:
                    break
                messages.extend(chunk)
                if len(chunk) < count:
                    break
                page += 1
            messages = messages[a.offset:a.offset + a.limit]
            if a.since:
                messages = [m for m in messages
                            if (m.get("receivedDate") or "") >= a.since]

        for m in messages:
            mid = str(m.get("id"))
            entry = {"id": mid, "subject": m.get("subject", ""),
                     "date": m.get("receivedDate", ""), "attachments": 0, "skipped": False}
            msg_dir = out_root / a.folder / mid
            md_path = msg_dir / "message.md"
            if md_path.exists() and not a.force:
                entry["skipped"] = True
                summary.append(entry)
                continue
            if a.dry_run:
                entry["skipped"] = "dry-run"
                summary.append(entry)
                continue
            full = OOCLIENT.get_message(mid)
            entry.update(convert_one(msg_dir, full, a.folder, out_root, a.ocr, a.no_attachments,
                                     target_dir=msg_dir))
            summary.append(entry)
    if a.json:
        print(json.dumps(summary, ensure_ascii=False, indent=2))
    else:
        imported = [e for e in summary if not e["skipped"]]
        print(f"mail/import: folder={a.folder} checked={len(summary)} "
              f"imported={len(imported)} (skipped={sum(e['skipped'] is True for e in summary)})")
        for e in summary:
            flag = "skip" if e["skipped"] is True else ("dry" if e["skipped"] == "dry-run" else "ok  ")
            print(f"  [{flag}] {e['id']} {e['date'][:10]} {e['subject'][:60]}"
                  f" (atts={e['attachments']})")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))