#!/usr/bin/env python3
"""
mail_archive.py — инкрементальный архиватор почты в локальную файловую базу.

Сохраняет письма из IMAP-ящика в структуру:
  /opt/hermes/email/<folder>/YYYY/MM/UID/
    ├── email.md          # YAML-frontmatter + тело письма
    └── attachments/      # вложения (если есть)

Формат email.md:
---
id: 255
folder: Sent
subject: "Re: ..."
from: "Name <addr>"
to: "Name <addr>"
date: "2025-07-10 11:51+03:00"
flags: ["Seen"]
has_attachment: false
message_id: <...@vinogorod.ru>
in_reply_to: <...@vinogorod.ru>
references: <...> <...>
cc: "Name <addr>"
content_type: multipart/mixed; boundary=...
---

Body text here...

Отслеживает last_uid для каждой папки в /opt/hermes/email/state/mail-archive-last-<folder>.json

Запуск:
  python3 ~/bin/mail_archive.py --folder INBOX --limit 100
  python3 ~/bin/mail_archive.py --folder "Отправленные" --limit 200
"""

import subprocess
import json
import sys
import argparse
import re
import hashlib
from datetime import datetime
from pathlib import Path

# Конфигурация
ARCHIVE_ROOT = Path("/opt/hermes/email")
STATE_DIR = ARCHIVE_ROOT / "state"
HIMALAYA_CMD = ["himalaya"]

# Папки для полной архивации (основные папки)
FOLDERS = ["INBOX", "Отправленные", "Archive", "Sent"]

# Вложенные папки INBOX, которые тоже архивируем
INBOX_SUBFOLDERS = [
    "INBOX/!Scan",
    "INBOX/!Битрикс",
    "INBOX/!ВГ Чек листы",
    "INBOX/!Документооборот",
    "INBOX/!Завки",
    "INBOX/!Материалы",
    "INBOX/!Отчеты",
    "INBOX/!Персонал",
    "INBOX/!Протоколы",
    "INBOX/!Реестр оплаты",
    "INBOX/!Торик",
    "INBOX/Бюджет",
    "INBOX/Контрагенты",
    "INBOX/ЛНД",
    "INBOX/Организация работы",
    "INBOX/Приемка и стройка",
    "INBOX/Системы",
    "INBOX/Эксплуатация",
]

# Какие дополнительные заголовки вытягивать через --header
EXTRA_HEADERS = [
    "Message-ID",
    "References",
    "In-Reply-To",
    "CC",
    "Content-Type",
]


def run_cmd(cmd, timeout=60):
    """Выполнить команду, вернуть stdout."""
    result = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        timeout=timeout,
    )
    if result.returncode != 0:
        stderr_text = result.stderr.decode("utf-8", errors="ignore")
        if "No such folder" in stderr_text:
            return ""
        raise RuntimeError(f"Command failed: {' '.join(cmd)}\n{stderr_text}")
    return result.stdout.decode("utf-8", errors="ignore")


def get_state_file(folder):
    """Получить путь к файлу состояния для папки."""
    safe_name = folder.replace("/", "_").replace(" ", "_")
    if not all(ord(c) < 128 for c in safe_name):
        h = hashlib.md5(folder.encode()).hexdigest()[:12]
        safe_name = f"folder_{h}"
    return STATE_DIR / f"mail-archive-last-{safe_name}.json"


def load_state(folder):
    """Загрузить last_uid для папки."""
    state_file = get_state_file(folder)
    if not state_file.exists():
        return {"last_uid": 0}
    try:
        with state_file.open("r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {"last_uid": 0}


def save_state(folder, state):
    """Сохранить last_uid для папки."""
    state_file = get_state_file(folder)
    state_file.parent.mkdir(parents=True, exist_ok=True)
    with state_file.open("w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


def parse_date(date_str):
    """Извлечь год/месяц из строки даты."""
    if not date_str:
        return datetime.now().year, datetime.now().month
    try:
        dt = datetime.fromisoformat(date_str)
        return dt.year, dt.month
    except Exception:
        pass
    try:
        dt = datetime.strptime(date_str[:25], "%a, %d %b %Y %H:%M:%S")
        return dt.year, dt.month
    except Exception:
        pass
    return datetime.now().year, datetime.now().month


def fmt_contact(c):
    """Форматировать контакт из {name, addr} в 'Name <addr>'."""
    if not c:
        return ""
    name = c.get("name") or ""
    addr = c.get("addr") or ""
    if name and addr:
        return f"{name} <{addr}>"
    return addr or name


def yaml_val(v):
    """Экранировать значение для YAML: строки в кавычки, если нужно."""
    return json.dumps(v, ensure_ascii=False)


def get_envelopes(folder, limit=100, last_uid=0):
    """
    Получить список конвертов из папки.
    Останавливается раньше, если набрал достаточно новых писем (UID > last_uid).
    """
    envelopes = []
    page = 1
    page_size = 500  # больше писем за страницу = меньше запросов

    while True:
        try:
            stdout = run_cmd(
                HIMALAYA_CMD + [
                    "envelope", "list",
                    "--folder", folder,
                    "--page", str(page),
                    "--page-size", str(page_size),
                    "--output", "json",
                ],
                timeout=60,
            )
        except RuntimeError as e:
            if "No such folder" in str(e):
                return []
            raise

        if not stdout.strip():
            break

        try:
            batch = json.loads(stdout)
        except json.JSONDecodeError:
            break

        if not batch:
            break

        envelopes.extend(batch)

        # Считаем сколько новых писем уже набрали
        new_count = 0
        for e in batch:
            uid = int(e.get("uid") or e.get("id") or 0)
            if uid > last_uid:
                new_count += 1

        # Хватит — не тащим остальные страницы
        total_new = sum(
            1 for e in envelopes
            if int(e.get("uid") or e.get("id") or 0) > last_uid
        )
        if total_new >= limit:
            break

        # Если ВСЕ письма на этой странице уже старые (UID <= last_uid) —
        # дальше можно не ходить, там только старее (сортировка по UID descending)
        if len(batch) < page_size:
            break
        last_batch_min = min(
            int(e.get("uid") or e.get("id") or 0) for e in batch
        )
        if last_batch_min <= last_uid:
            break

        page += 1

    return envelopes


def get_email_content(uid, folder):
    """
    Получить тело письма и дополнительные заголовки.
    Возвращает (headers_dict, body_text).
    """
    # Собираем аргументы --header
    header_args = []
    for h in EXTRA_HEADERS:
        header_args.extend(["--header", h])

    try:
        stdout = run_cmd(
            HIMALAYA_CMD + [
                "message", "read", str(uid),
                "--folder", folder,
            ] + header_args,
            timeout=30,
        )
    except RuntimeError:
        return {}, "(body unavailable)\n"

    # Разделяем на блок заголовков и тело
    # Первый \n\n — граница между запрошенными заголовками и остальным
    parts = stdout.split("\n\n", 1)
    header_block = parts[0] if parts else ""
    body = parts[1] if len(parts) > 1 else ""

    # Парсим заголовки
    headers = {}
    for line in header_block.split("\n"):
        if ":" in line:
            key, val = line.split(":", 1)
            headers[key.strip()] = val.strip()

    return headers, body


def make_email_md(meta, extra_headers, body, folder_name):
    """
    Собрать email.md с YAML-frontmatter из meta.json + extra_headers + body.
    """
    lines = []
    lines.append("---")

    # Основные поля из envelope
    lines.append(f"id: {meta.get('id', '0')}")
    lines.append(f"folder: {folder_name}")
    lines.append(f"subject: {yaml_val(meta.get('subject', ''))}")

    from_str = fmt_contact(meta.get("from"))
    to_str = fmt_contact(meta.get("to"))
    lines.append(f"from: {yaml_val(from_str)}")
    lines.append(f"to: {yaml_val(to_str)}")

    date_str = meta.get("date") or meta.get("internal_date") or ""
    lines.append(f"date: {yaml_val(date_str)}")
    lines.append(f"flags: {json.dumps(meta.get('flags', []))}")
    lines.append(f"has_attachment: {str(meta.get('has_attachment', False)).lower()}")

    # Дополнительные заголовки
    for h in EXTRA_HEADERS:
        val = extra_headers.get(h)
        if val:
            lines.append(f"{h}: {yaml_val(val)}")

    lines.append("---")
    lines.append("")

    # Тело письма
    lines.append(body.rstrip("\n"))

    return "\n".join(lines)


def get_attachments(uid, folder, dest_dir):
    """Скачать вложения письма в dest_dir."""
    try:
        run_cmd(
            HIMALAYA_CMD + [
                "attachment", "download", str(uid),
                "--folder", folder,
                "--dir", str(dest_dir),
            ],
            timeout=60,
        )
    except RuntimeError:
        pass  # нет вложений — норм


def archive_folder(folder, limit=100):
    """
    Архивировать все новые письма из папки.
    Возвращает количество заархивированных писем.
    """
    state = load_state(folder)
    last_uid = state.get("last_uid", 0)

    print(f"\n📁 {folder} (last_uid={last_uid})")

    envelopes = get_envelopes(folder, limit=limit, last_uid=last_uid)
    if not envelopes:
        print(f"  → нет писем или папка не найдена")
        return 0

    # Отфильтровать новые по UID
    new = []
    for env in envelopes:
        uid = env.get("uid") or env.get("id") or 0
        try:
            uid = int(uid)
        except (ValueError, TypeError):
            continue
        if uid > last_uid:
            new.append((uid, env))

    if not new:
        print(f"  → новых писем нет")
        return 0

    new.sort(key=lambda x: x[0])

    processed = 0
    for uid, env in new:
        if processed >= limit:
            print(f"  → достигнут лимит {limit}, осталось ещё {len(new) - processed}")
            break

        # Год/месяц
        date_str = env.get("date") or env.get("internal_date") or ""
        year, month = parse_date(date_str)

        # Путь: /mnt/yandex-disk/hermes/email/<folder>/YYYY/MM/UID/
        msg_dir = ARCHIVE_ROOT / folder / f"{year:04d}" / f"{month:02d}" / str(uid)
        email_path = msg_dir / "email.md"
        attachments_dir = msg_dir / "attachments"

        # Проверка — уже сохранено
        if email_path.exists():
            print(f"  UID {uid} уже есть, пропускаю")
            last_uid = max(last_uid, uid)
            processed += 1
            continue

        msg_dir.mkdir(parents=True, exist_ok=True)
        attachments_dir.mkdir(parents=True, exist_ok=True)

        # Получаем заголовки и тело
        extra_headers, body = get_email_content(uid, folder)

        # Собираем email.md
        content = make_email_md(env, extra_headers, body, folder)
        email_path.write_text(content, encoding="utf-8")

        # Вложения
        get_attachments(uid, folder, attachments_dir)

        subj = (env.get("subject") or "")[:60]
        print(f"  ✓ UID {uid} ({subj})")
        last_uid = uid
        processed += 1

    # Сохранить состояние
    state["last_uid"] = last_uid
    save_state(folder, state)
    print(f"  → last_uid обновлён до {last_uid}, обработано {processed}")
    return processed


def main():
    parser = argparse.ArgumentParser(
        description="Архиватор почты — инкрементальное сохранение в /mnt/yandex-disk/hermes/email/"
    )
    parser.add_argument(
        "--folder", default=None,
        help="Архивировать только одну папку (по умолчанию: все основные папки)"
    )
    parser.add_argument(
        "--limit", type=int, default=200,
        help="Максимум писем за один запуск (по умолчанию: 200)"
    )
    parser.add_argument(
        "--all", action="store_true",
        help="Архивировать включая вложенные папки INBOX"
    )
    args = parser.parse_args()

    # Определить список папок
    if args.folder:
        folders_to_archive = [args.folder]
    elif args.all:
        folders_to_archive = FOLDERS + INBOX_SUBFOLDERS
    else:
        folders_to_archive = FOLDERS

    total = 0
    for folder in folders_to_archive:
        try:
            count = archive_folder(folder, limit=args.limit)
            total += count
        except Exception as e:
            print(f"  [ERROR] {folder}: {e}", file=sys.stderr)

    print(f"\n{'='*50}")
    print(f"Готово. Всего заархивировано писем: {total}")
    return 0


if __name__ == "__main__":
    sys.exit(main())