#!/usr/bin/env python3
"""IMAP IDLE daemon for routing new mail through mailcat rules."""

from __future__ import annotations

import argparse
import imaplib
import json
import logging
import select
import signal
import sys
import time
from pathlib import Path

import config
from imap_utils import quote_mailbox
from mail_imap_ops import HEADER_FETCH, destination_from_header, ensure_mailbox, fetch_body, move_message

DEFAULT_STATE_FILE = Path(__file__).parent / "sort_mail_daemon_state.json"
DEFAULT_LOG_FILE = Path(__file__).parent / "sort_mail_daemon.log"


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Sort new IMAP mail using mailcat routes.")
    parser.add_argument("--account", default="backup", help="config.json account name")
    parser.add_argument("--folder", default="INBOX", help="folder to watch")
    parser.add_argument("--state-file", type=Path, default=DEFAULT_STATE_FILE)
    parser.add_argument("--log-file", type=Path, default=DEFAULT_LOG_FILE)
    parser.add_argument("--idle-timeout", type=int, default=1740)
    parser.add_argument("--reconnect-delay", type=int, default=30)
    parser.add_argument("--once", action="store_true", help="process once and exit")
    parser.add_argument("--dry-run", action="store_true", help="log intended moves without changing IMAP")
    return parser


def load_state(path: Path) -> dict:
    if not path.exists():
        return {"folder_state": {}}
    try:
        state = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {"folder_state": {}}
    state.setdefault("folder_state", {})
    return state


def save_state(path: Path, state: dict):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(state, indent=2, ensure_ascii=False), encoding="utf-8")


def folder_state(state: dict, folder: str) -> dict:
    folders = state.setdefault("folder_state", {})
    entry = folders.setdefault(folder, {})
    entry.setdefault("uidvalidity", None)
    entry.setdefault("processed_uids", [])
    return entry


def reset_on_uidvalidity_change(state: dict, folder: str, uidvalidity: str | None):
    entry = folder_state(state, folder)
    if entry.get("uidvalidity") != uidvalidity:
        entry["uidvalidity"] = uidvalidity
        entry["processed_uids"] = []


def processed_set(state: dict, folder: str) -> set[str]:
    return set(folder_state(state, folder).get("processed_uids", []))


def mark_processed(state: dict, folder: str, uid: str):
    entry = folder_state(state, folder)
    processed = set(entry.get("processed_uids", []))
    processed.add(uid)
    entry["processed_uids"] = sorted(processed, key=lambda value: int(value) if value.isdigit() else value)


def setup_logging(log_file: Path):
    log_file.parent.mkdir(parents=True, exist_ok=True)
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
        handlers=[logging.FileHandler(log_file, encoding="utf-8"), logging.StreamHandler(sys.stdout)],
    )


def connect(account_name: str):
    cfg = config.load(account_name)
    mail = imaplib.IMAP4_SSL(cfg.imap_host, cfg.imap_port)
    mail.login(cfg.username, cfg.password)
    logging.info("connected account=%s username=%s host=%s", cfg.name, cfg.username, cfg.imap_host)
    return mail


def select_folder(mail, folder: str) -> str | None:
    status, data = mail.select(quote_mailbox(folder))
    if status != "OK":
        raise RuntimeError(f"Cannot select {folder}: {data!r}")
    _status, data = mail.response("UIDVALIDITY")
    if data and data[0]:
        return data[0].decode(errors="replace") if isinstance(data[0], bytes) else str(data[0])
    return None


def uid_search_all(mail) -> list[bytes]:
    status, data = mail.uid("SEARCH", None, "ALL")
    if status != "OK" or not data or not data[0]:
        return []
    return [uid for uid in data[0].split() if uid]


def process_folder(mail, state: dict, folder: str, dry_run: bool) -> dict[str, int]:
    uidvalidity = select_folder(mail, folder)
    reset_on_uidvalidity_change(state, folder, uidvalidity)
    processed = processed_set(state, folder)
    counts = {"seen": 0, "matched": 0, "moved": 0, "unmatched": 0, "failed": 0, "skipped": 0}

    for uid in uid_search_all(mail):
        uid_text = uid.decode(errors="replace")
        counts["seen"] += 1
        if uid_text in processed:
            counts["skipped"] += 1
            continue

        status, data = mail.uid("FETCH", uid, HEADER_FETCH)
        raw_header = fetch_body(data)
        if status != "OK" or raw_header is None:
            counts["failed"] += 1
            logging.warning("fetch_failed folder=%s uid=%s status=%s data=%r", folder, uid_text, status, data)
            continue

        destination = destination_from_header(raw_header)
        if not destination or destination == folder:
            counts["unmatched"] += 1
            mark_processed(state, folder, uid_text)
            logging.info("no_match folder=%s uid=%s", folder, uid_text)
            continue

        counts["matched"] += 1
        if dry_run:
            logging.info("dry_move folder=%s uid=%s destination=%s", folder, uid_text, destination)
            mark_processed(state, folder, uid_text)
            continue

        if not ensure_mailbox(mail, destination):
            counts["failed"] += 1
            logging.error("ensure_destination_failed folder=%s uid=%s destination=%s", folder, uid_text, destination)
            continue

        ok, failure = move_message(mail, uid, folder, destination)
        if ok:
            counts["moved"] += 1
            mark_processed(state, folder, uid_text)
            logging.info("moved folder=%s uid=%s destination=%s", folder, uid_text, destination)
        else:
            counts["failed"] += 1
            logging.error("move_failed folder=%s uid=%s destination=%s failure=%r", folder, uid_text, destination, failure)

    if not dry_run:
        mail.expunge()
    return counts


def wait_for_mail(mail, timeout: int) -> bool:
    if hasattr(mail, "idle"):
        with mail.idle(duration=timeout) as idler:
            for typ, _data in idler:
                if typ == "EXISTS":
                    return True
        return False
    return wait_for_mail_private_idle(mail, timeout)


def wait_for_mail_private_idle(mail, timeout: int) -> bool:
    tag = mail._new_tag()
    mail.send(tag + b" IDLE\r\n")
    continuation = mail.readline()
    if not continuation.startswith(b"+"):
        raise RuntimeError(f"IDLE rejected: {continuation!r}")

    changed = False
    try:
        ready, _, _ = select.select([mail.sock], [], [], timeout)
        if ready:
            line = mail.readline()
            changed = b"EXISTS" in line or b"RECENT" in line
    finally:
        mail.send(b"DONE\r\n")
        while True:
            line = mail.readline()
            if line.startswith(tag):
                break
    return changed


def run(args) -> int:
    setup_logging(args.log_file)
    stop = {"requested": False}

    def request_stop(_signum, _frame):
        stop["requested"] = True

    signal.signal(signal.SIGTERM, request_stop)
    signal.signal(signal.SIGINT, request_stop)

    state = load_state(args.state_file)

    while not stop["requested"]:
        mail = None
        try:
            mail = connect(args.account)
            counts = process_folder(mail, state, args.folder, args.dry_run)
            save_state(args.state_file, state)
            logging.info("scan_complete folder=%s counts=%s", args.folder, counts)
            if args.once:
                return 0

            while not stop["requested"]:
                select_folder(mail, args.folder)
                changed = wait_for_mail(mail, args.idle_timeout)
                if changed:
                    counts = process_folder(mail, state, args.folder, args.dry_run)
                    save_state(args.state_file, state)
                    logging.info("idle_scan_complete folder=%s counts=%s", args.folder, counts)
                else:
                    logging.info("idle_timeout folder=%s", args.folder)
        except Exception:
            logging.exception("daemon_error reconnecting_after=%ss", args.reconnect_delay)
            if args.once:
                return 1
            time.sleep(args.reconnect_delay)
        finally:
            if mail is not None:
                try:
                    mail.logout()
                except Exception:
                    pass

    save_state(args.state_file, state)
    logging.info("stopped")
    return 0


if __name__ == "__main__":
    raise SystemExit(run(build_parser().parse_args()))
