#!/usr/bin/env python3
"""
Gmail IMAP email exporter v3 — optimized.
Single IMAP query for all emails since Dec 20, 2025, then filter by domains in Python.
Each email becomes a PDF with body + attachments merged.
Organized by domain and also in a flat 'all' folder.
"""

import os
import sys
import re
import time
import logging
from pathlib import Path
from datetime import datetime, date
from email.utils import parsedate_to_datetime

from imap_tools import MailBox, AND
from weasyprint import HTML
from fpdf import FPDF
import fitz  # PyMuPDF

# === CONFIG ===
IMAP_SERVER = "imap.gmail.com"

# Извлекаем учётку из .muttrc (сначала в проекте, потом в home)
MUTT_CONF = Path(__file__).parent / ".muttrc"
if not MUTT_CONF.exists():
    MUTT_CONF = Path.home() / ".muttrc"

# Domains to search
DOMAINS = [
    "gosuslugi.ru", "dom.gosuslugi.ru", "gov.ru", "duma.gov.ru",
    "mvd.gov.ru", "mvd.ru", "rkn.gov.ru", "nalog.gov.ru", "nalog.ru",
    "genproc.gov.ru", "pfr.gov.ru", "fss.ru", "rospotrebnadzor.ru",
    "russianpost.ru", "mailop.ru", "73.mailop.ru", "uk-service.ru",
    "ulgss.ru", "ulgov.ru", "ulgkh.ru", "dgi.ru", "mos.ru", "pfrf.ru",
    "fsin.gov.ru", "skrf.ru", "epp.genproc.gov.ru", "ugpr.ru",
    "fas.gov.ru", "minstroyrf.gov.ru", "gkh.ru", "sudrf.ru",
    "msud.ru", "arbitr.ru", "fssp.gov.ru", "fssprus.ru", "mirsud",
    "msud73.ru", "mirsud73.ru", "court.ru", "sud.ru", "vkks.ru",
    "uloblsud.ru", "sledcom.ru"
]

# From this date
SINCE_DATE = date(2025, 12, 20)

# Gmail folder to search
GMAIL_FOLDER = "[Gmail]/Вся почта"

# Output
BASE_OUTPUT = Path(__file__).parent / "exported_emails"
ALL_FOLDER = BASE_OUTPUT / "all"

# Progress tracking
PROGRESS_FILE = Path(__file__).parent / ".export_progress"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(Path(__file__).parent / "export.log", encoding="utf-8"),
        logging.StreamHandler(sys.stdout),
    ],
)
log = logging.getLogger(__name__)


def load_progress():
    """Load set of already processed UIDs."""
    if PROGRESS_FILE.exists():
        return set(PROGRESS_FILE.read_text().strip().splitlines())
    return set()


def save_progress(uid, processed_set):
    """Save progress after each email."""
    processed_set.add(str(uid))
    PROGRESS_FILE.write_text("\n".join(processed_set))


def sanitize_filename(name, max_len=80):
    """Make filename safe."""
    name = re.sub(r'[<>:"/\\|?*\x00-\x1f]', '_', name)
    name = re.sub(r'\s+', ' ', name).strip()
    return name[:max_len]


def parse_mutt_config():
    """Extract IMAP user and password from .muttrc."""
    config = {}
    with open(MUTT_CONF) as f:
        for line in f:
            line = line.strip()
            if line.startswith('#') or not line:
                continue
            match = re.match(r'set\s+(\w+)\s*=\s*"([^"]*)"', line)
            if not match:
                match = re.match(r'set\s+(\w+)\s*=\s*(\S+)', line)
            if match:
                key, val = match.group(1), match.group(2)
                config[key] = val
    return config


def html_to_pdf_stream(html_content):
    """Convert HTML to PDF bytes using weasyprint, fallback to fpdf."""
    try:
        html_doc = HTML(string=html_content)
        return html_doc.write_pdf()
    except Exception as e:
        log.warning(f"weasyprint failed, using fpdf fallback: {e}")
        return html_to_pdf_fpdf(html_content)


def html_to_pdf_fpdf(html_content):
    """Fallback: strip HTML and use FPDF."""
    from bs4 import BeautifulSoup
    soup = BeautifulSoup(html_content, 'lxml')
    text = soup.get_text(separator='\n')

    pdf = FPDF()
    pdf.add_page()
    pdf.add_font('DejaVu', '', '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf', uni=True)
    pdf.add_font('DejaVu', 'B', '/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf', uni=True)
    pdf.set_font('DejaVu', '', 10)

    lines = [line for line in text.split('\n') if line.strip()]
    for line in lines:
        pdf.multi_cell(0, 6, line)
        pdf.ln(1)

    return pdf.output()


def convert_attachment_to_pdf(attachment_data, filename, content_type):
    """Convert an attachment to PDF."""
    ext = Path(filename).suffix.lower()

    try:
        if ext in ('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp', '.svg'):
            pdf = FPDF()
            pdf.add_page()
            img_w, img_h = get_image_dims(attachment_data)
            page_w = pdf.w - 20
            scale = min(page_w / img_w, (pdf.h - 20) / img_h, 1.0)
            w, h = img_w * scale, img_h * scale
            import tempfile
            with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp:
                tmp.write(attachment_data)
                tmp_path = tmp.name
            try:
                pdf.image(tmp_path, x=10, y=10, w=w, h=h)
            finally:
                os.unlink(tmp_path)
            return pdf.output()

        elif ext in ('.doc', '.docx'):
            import tempfile, subprocess
            with tempfile.TemporaryDirectory() as tmpdir:
                src = Path(tmpdir) / filename
                src.write_bytes(attachment_data)
                subprocess.run(
                    ['libreoffice', '--headless', '--convert-to', 'pdf',
                     '--outdir', tmpdir, str(src)],
                    capture_output=True, timeout=120
                )
                pdf_path = src.with_suffix('.pdf')
                if pdf_path.exists():
                    return pdf_path.read_bytes()

        elif ext == '.pdf':
            return attachment_data

        elif ext in ('.xls', '.xlsx'):
            import tempfile, subprocess
            with tempfile.TemporaryDirectory() as tmpdir:
                src = Path(tmpdir) / filename
                src.write_bytes(attachment_data)
                subprocess.run(
                    ['libreoffice', '--headless', '--convert-to', 'pdf',
                     '--outdir', tmpdir, str(src)],
                    capture_output=True, timeout=120
                )
                pdf_path = src.with_suffix('.pdf')
                if pdf_path.exists():
                    return pdf_path.read_bytes()

        elif ext in ('.txt', '.csv', '.log', '.xml', '.html', '.htm'):
            try:
                text = attachment_data.decode('utf-8', errors='replace')
            except:
                text = attachment_data.decode('latin-1', errors='replace')
            return html_to_pdf_fpdf(text)

        else:
            info = f"Attachment: {filename}\nType: {content_type}\nSize: {len(attachment_data)} bytes\n\n[Binary content - not converted]"
            return html_to_pdf_fpdf(info)

    except Exception as e:
        log.warning(f"Failed to convert {filename}: {e}")
        info = f"Attachment: {filename}\nType: {content_type}\nSize: {len(attachment_data)} bytes\n\nConversion failed: {e}"
        return html_to_pdf_fpdf(info)


def get_image_dims(data):
    """Get image dimensions."""
    import io
    from PIL import Image
    img = Image.open(io.BytesIO(data))
    w_px, h_px = img.size
    w_mm = w_px * 25.4 / 96
    h_mm = h_px * 25.4 / 96
    return w_mm, h_mm


def merge_pdfs(pdf_bytes_list):
    """Merge multiple PDFs into one."""
    if not pdf_bytes_list:
        return None
    if len(pdf_bytes_list) == 1:
        return pdf_bytes_list[0]

    merger = fitz.open()
    for pdf_bytes in pdf_bytes_list:
        try:
            doc = fitz.open(stream=pdf_bytes, filetype="pdf")
            merger.insert_pdf(doc)
        except Exception as e:
            log.warning(f"Failed to merge a PDF section: {e}")

    output = merger.tobytes()
    merger.close()
    return output


def build_email_html(msg):
    """Build nice HTML representation of email with headers."""
    subject = msg.subject or "(Без темы)"
    from_addr = msg.from_ or "unknown"
    to_addr = ", ".join(msg.to) if msg.to else "unknown"
    cc_addr = ", ".join(msg.cc) if msg.cc else ""
    date_str = msg.date_str or str(msg.date) if msg.date else "unknown"

    body_html = msg.html
    if not body_html and msg.text:
        text = msg.text
        text_escaped = text.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
        body_html = text_escaped.replace('\n', '<br>\n')
    elif not body_html:
        body_html = "<p><i>(Нет тела письма)</i></p>"

    attachments_info = ""
    if msg.attachments:
        att_items = []
        for att in msg.attachments:
            att_items.append(f'<li>{att.filename} ({att.content_type}, {att.size} bytes)</li>')
        attachments_info = '<h3 style="color:#555;">Вложения:</h3><ul>' + ''.join(att_items) + '</ul>'

    html = f"""
    <html>
    <head>
    <meta charset="utf-8">
    <style>
        body {{ font-family: 'DejaVu Sans', Arial, sans-serif; font-size: 12px; margin: 20px; color: #333; }}
        .header {{ background: #f0f4f8; padding: 15px; border-radius: 5px; margin-bottom: 15px; border-left: 4px solid #3b82f6; }}
        .header h2 {{ margin: 0 0 10px 0; font-size: 16px; color: #1e40af; }}
        .header table {{ width: 100%; }}
        .header td {{ padding: 3px 5px; vertical-align: top; }}
        .header td:first-child {{ font-weight: bold; width: 80px; color: #555; }}
        .body {{ line-height: 1.6; }}
        .attachments {{ margin-top: 20px; padding: 10px; background: #f9fafb; border-radius: 3px; }}
        .attachments li {{ margin: 3px 0; color: #666; }}
    </style>
    </head>
    <body>
    <div class="header">
        <h2>{subject}</h2>
        <table>
            <tr><td>От:</td><td>{from_addr}</td></tr>
            <tr><td>Кому:</td><td>{to_addr}</td></tr>
            {f'<tr><td>Копия:</td><td>{cc_addr}</td></tr>' if cc_addr else ''}
            <tr><td>Дата:</td><td>{date_str}</td></tr>
        </table>
    </div>
    <div class="body">
        {body_html}
    </div>
    {f'<div class="attachments">{attachments_info}</div>' if attachments_info else ''}
    </body>
    </html>
    """
    return html


def matches_domain(addr, domains):
    """Check if address matches any domain. Returns matched domain or None."""
    if not addr:
        return None
    addr_lower = addr.lower()
    for domain in domains:
        if domain.lower() in addr_lower:
            return domain
    return None


def main():
    log.info("=" * 60)
    log.info("Gmail Email Exporter v3 — starting")
    log.info("=" * 60)

    # Parse config
    config = parse_mutt_config()
    imap_user = config.get('imap_user', '')
    imap_pass = config.get('imap_pass', '')

    if not imap_user or imap_user == 'your.email@gmail.com':
        log.error("ERROR: Не настроены учётные данные в .muttrc!")
        log.error("Укажите реальный imap_user и imap_pass (App Password)")
        sys.exit(1)

    log.info(f"IMAP User: {imap_user}")
    log.info(f"Domains: {len(DOMAINS)} domains")
    log.info(f"Since: {SINCE_DATE.strftime('%Y-%m-%d')}")

    # Create output dirs
    BASE_OUTPUT.mkdir(parents=True, exist_ok=True)
    ALL_FOLDER.mkdir(parents=True, exist_ok=True)

    # Load progress
    processed = load_progress()
    log.info(f"Previously processed: {len(processed)} emails")

    # Connect
    log.info("Connecting to Gmail IMAP...")
    try:
        mailbox = MailBox(IMAP_SERVER).login(imap_user, imap_pass)
        log.info("Connected successfully!")
    except Exception as e:
        log.error(f"Failed to connect: {e}")
        sys.exit(1)

    # Auto-detect folder names
    all_mail_folder = "[Gmail]/Вся почта"
    spam_folder = "[Gmail]/Спам"
    try:
        for folder in mailbox.folder.list():
            if r'\All' in folder.flags:
                all_mail_folder = folder.name
            elif r'\Junk' in folder.flags:
                spam_folder = folder.name
    except Exception as e:
        log.warning(f"Failed to auto-detect folder names, using defaults: {e}")

    folders_to_scan = [all_mail_folder, spam_folder]
    matching_uids = []
    total_scanned = 0

    # Phase 1: Fetch headers from all target folders
    log.info("Phase 1: Scanning emails (headers only) to find matches...")
    scan_start = time.time()
    last_log_time = time.time()

    for folder_name in folders_to_scan:
        log.info(f"Scanning folder: {folder_name}")
        try:
            mailbox.folder.set(folder_name)
        except Exception as e:
            log.error(f"Failed to set folder {folder_name}: {e}")
            continue

        for msg in mailbox.fetch(
            AND(date_gte=SINCE_DATE),
            bulk=True,
            headers_only=True,
            mark_seen=False,
        ):
            total_scanned += 1
            now = time.time()
            if now - last_log_time > 10:
                elapsed = now - scan_start
                rate = total_scanned / elapsed if elapsed > 0 else 0
                log.info(f"  Scanned in {folder_name}: {total_scanned} ({rate:.0f}/sec), Matching so far: {len(matching_uids)}")
                last_log_time = now

            # Check if matches any domain
            from_match = matches_domain(msg.from_, DOMAINS)
            to_match = None
            if not from_match:
                for to_addr in (msg.to or []):
                    to_match = matches_domain(to_addr, DOMAINS)
                    if to_match:
                        break

            if from_match or to_match:
                msg_date = msg.date
                if msg_date and hasattr(msg_date, 'date'):
                    msg_date = msg_date.date()
                if not msg_date or msg_date >= SINCE_DATE:
                    matching_uids.append((msg.uid, folder_name))

    scan_elapsed = time.time() - scan_start
    log.info(f"Phase 1 complete: scanned {total_scanned} emails in {scan_elapsed/60:.1f}min")
    log.info(f"Found {len(matching_uids)} matching emails across folders")

    if not matching_uids:
        log.info("No matching emails found. Exiting.")
        mailbox.logout()
        return

    # Phase 2: Download full emails and convert to PDF
    log.info("\n" + "=" * 60)
    log.info("Phase 2: Downloading full emails and converting to PDF...")
    log.info("=" * 60)

    start_time = time.time()
    total_processed = 0
    total_errors = 0

    current_folder = ""

    def reconnect_mailbox(target_folder):
        """Reconnect to IMAP with retry and set folder."""
        for attempt in range(5):
            try:
                log.info(f"Reconnecting... attempt {attempt+1}")
                mb = MailBox(IMAP_SERVER).login(imap_user, imap_pass)
                mb.folder.set(target_folder)
                log.info(f"Reconnected and folder set to {target_folder}!")
                return mb
            except Exception as e:
                log.warning(f"Reconnect failed: {e}")
                time.sleep(10 * (attempt + 1))
        log.error("Failed to reconnect after 5 attempts")
        sys.exit(1)

    for idx, (uid, f_name) in enumerate(matching_uids, 1):
        if current_folder != f_name:
            try:
                mailbox.folder.set(f_name)
                current_folder = f_name
            except Exception as e:
                log.error(f"Failed to switch to folder {f_name}: {e}")
                continue

        progress_key = f"{f_name}:{uid}" if f_name == spam_folder else str(uid)
        if progress_key in processed or str(uid) in processed:
            log.info(f"[{idx}/{len(matching_uids)}] SKIP (already processed) UID={uid} in {f_name}")
            continue

        elapsed = time.time() - start_time
        eta = (elapsed / (idx - 1)) * (len(matching_uids) - idx + 1) if idx > 1 else 0

        log.info(f"[{idx}/{len(matching_uids)}] Downloading UID={uid} from {f_name} | ETA: {eta/60:.1f}m")

        try:
            msgs = list(mailbox.fetch(
                AND(uid=uid),
                headers_only=False,
                mark_seen=False,
            ))
            if not msgs:
                log.warning(f"  Could not fetch UID={uid} in {f_name}")
                total_errors += 1
                save_progress(progress_key, processed)
                continue

            msg = msgs[0]

            # Build email info
            subject = msg.subject or "(Без темы)"
            from_addr = msg.from_ or "unknown"
            date_str = msg.date_str or "unknown"

            # Determine matching domain
            match_domain = matches_domain(from_addr, DOMAINS)
            if not match_domain:
                for to_addr in (msg.to or []):
                    match_domain = matches_domain(to_addr, DOMAINS)
                    if match_domain:
                        break
            if not match_domain:
                match_domain = "other"

            safe_subject = sanitize_filename(subject)
            safe_from = sanitize_filename(from_addr.split('<')[-1].rstrip('>').split('@')[0], 30)
            date_part = msg.date.strftime('%Y-%m-%d') if msg.date else 'unknown'

            filename = f"{date_part}_{safe_from}_{safe_subject}.pdf"
            if len(filename) > 150:
                filename = filename[:150] + ".pdf"

            # Build HTML for email body
            html_content = build_email_html(msg)

            pdf_parts = []

            # Convert body to PDF
            body_pdf = html_to_pdf_stream(html_content)
            if body_pdf:
                pdf_parts.append(body_pdf)

            # Convert attachments to PDF
            if msg.attachments:
                log.info(f"  Has {len(msg.attachments)} attachment(s)")
                for att in msg.attachments:
                    try:
                        att_pdf = convert_attachment_to_pdf(
                            att.payload, att.filename, att.content_type
                        )
                        if att_pdf:
                            pdf_parts.append(att_pdf)
                    except Exception as e:
                        log.warning(f"  Failed to convert attachment {att.filename}: {e}")

            # Merge all PDFs
            if pdf_parts:
                merged_pdf = merge_pdfs(pdf_parts)

                if merged_pdf:
                    # Save to domain folder
                    domain_folder = BASE_OUTPUT / sanitize_filename(match_domain, 50)
                    domain_folder.mkdir(parents=True, exist_ok=True)

                    filepath_domain = domain_folder / filename
                    filepath_domain.write_bytes(merged_pdf)

                    # Save to all folder
                    filepath_all = ALL_FOLDER / f"{match_domain}_{filename}"
                    filepath_all.write_bytes(merged_pdf)

                    log.info(f"  SAVED: {filename}")
                    total_processed += 1
                else:
                    log.warning(f"  Failed to create merged PDF")
                    total_errors += 1
            else:
                log.warning(f"  No content to convert")
                total_errors += 1

            # Save progress
            save_progress(progress_key, processed)

        except Exception as e:
            # Check if it's a socket error — try to reconnect
            if 'socket error' in str(e) or 'EOF' in str(e) or 'abort' in str(e):
                log.warning(f"Connection lost, trying to reconnect...")
                mailbox = reconnect_mailbox(f_name)
                current_folder = f_name
                total_errors += 1
                continue
            else:
                log.error(f"  ERROR processing UID={uid} in {f_name}: {e}", exc_info=True)
                total_errors += 1
                save_progress(progress_key, processed)

    mailbox.logout()

    # Summary
    total_time = time.time() - start_time
    summary_text = f"""
📦 <b>ОТЧЕТ ПО ЭКСПОРТУ ПОЧТЫ (Mutt)</b>

📅 <b>Дата:</b> {datetime.now().strftime('%Y-%m-%d %H:%M')}
📧 <b>Просканировано:</b> {total_scanned}
🎯 <b>Найдено (фильтр):</b> {len(matching_uids)}
✅ <b>Обработано (новых):</b> {total_processed}
❌ <b>Ошибок:</b> {total_errors}
⏱️ <b>Время:</b> {total_time/60:.1f} мин.

📁 <i>Все файлы сохранены в PDF и готовы к загрузке в NotebookLM!</i>
"""
    log.info("\n" + "=" * 60)
    log.info("EXPORT COMPLETE!")
    log.info(f"Total scanned: {total_scanned}")
    log.info(f"Matching: {len(matching_uids)}")
    log.info(f"Processed this run: {total_processed}")
    log.info(f"Errors: {total_errors}")
    log.info(f"Total time: {total_time/60:.1f} minutes")
    log.info("=" * 60)

    # Telegram Notification (only if processed > 0)
    if total_processed > 0:
        try:
            import requests
            BOT_TOKEN = "8725618164:AAH1tGalq-pw1l0t4P5c0sdLkCVJdh7IE7M"
            CHAT_ID = "197957361"
            url = f"https://api.telegram.org/bot{BOT_TOKEN}/sendMessage"
            params = {
                "chat_id": CHAT_ID,
                "text": "📬 <b>ВАЖНОЕ ОБНОВЛЕНИЕ ПОЧТЫ!</b>\n" + summary_text + "\n🚀 <i>Пора обновить источники в NotebookLM!</i>",
                "parse_mode": "HTML"
            }
            # Send directly (NanoPi routes to VPN transparently)
            requests.post(url, json=params, timeout=20)
            log.info("✅ Telegram notification sent.")
        except Exception as e:
            log.error(f"❌ Failed to send Telegram notification: {e}")

if __name__ == '__main__':
    main()