#!/usr/bin/env python3
"""
Gmail Email Auditor v2.6 SPEEDSTER for Domovoy Bot.
1. Fast header-only scan to find matching UIDs.
2. Full download only for matches (with attachments).
3. Merges all PDFs into a master archive for NotebookLM.
4. Updates domovoy.db and notifies via Telegram.
"""

import os
import sys
import re
import time
import logging
import sqlite3
import tempfile
import subprocess
from pathlib import Path
from datetime import datetime, date, timedelta, timezone
import requests

from imap_tools import MailBox, AND
from weasyprint import HTML
import fitz  # PyMuPDF
from PIL import Image
from bs4 import BeautifulSoup
from dotenv import load_dotenv

# === CONFIG & PATHS ===
BASE_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = BASE_DIR.parent.parent
load_dotenv(PROJECT_ROOT / ".env")

DB_PATH = PROJECT_ROOT / "database" / "domovoy.db"
MUTT_CONF = BASE_DIR / ".muttrc"
EXPORT_DIR = BASE_DIR / "exported_emails"
ALL_FOLDER = EXPORT_DIR / "all"
MASTER_PDF = EXPORT_DIR / "TOTAL_ARCHIVE_2025-2026.pdf"
PROGRESS_FILE = BASE_DIR / ".export_progress"

IMAP_SERVER = "imap.gmail.com"
SINCE_DATE = date(2025, 12, 20)

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"
]

KEYWORDS = ["суд", "судье", "судебн", "дело №", "иск", "заседание"]

# Bot notification config
BOT_TOKEN = os.getenv("BOT_TOKEN")
ADMIN_USER_ID = os.getenv("ADMIN_USER_ID")

# Logging configuration with Ulyanovsk timezone (UTC+4)
ULY_TZ_OFFSET = 4 * 3600

class UlyanovskFormatter(logging.Formatter):
    def converter(self, timestamp):
        return datetime.fromtimestamp(timestamp, tz=timezone(timedelta(hours=4))).timetuple()
    def formatTime(self, record, datefmt=None):
        dt = datetime.fromtimestamp(record.created, tz=timezone(timedelta(hours=4)))
        return dt.strftime(datefmt) if datefmt else dt.isoformat()

log_formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(message)s')
# We'll use a simpler way - just set the converter for the formatter
log_formatter.converter = lambda *args: (datetime.now(timezone(timedelta(hours=4))).timetuple())

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler(BASE_DIR / "auditor_run.log", encoding="utf-8"),
        logging.StreamHandler(sys.stdout),
    ],
)
# Force all handlers to use Ulyanovsk time in their timestamps
for handler in logging.root.handlers:
    handler.formatter.converter = lambda *args: (datetime.now(timezone(timedelta(hours=4))).timetuple())

log = logging.getLogger("EmailAuditor")

# === UTILS ===

def parse_mutt_config():
    config = {}
    if not MUTT_CONF.exists(): return config
    with open(MUTT_CONF) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'): 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: config[match.group(1)] = match.group(2)
    return config

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

def send_bot_message(text):
    if not BOT_TOKEN or not ADMIN_USER_ID: return
    url = f"https://api.telegram.org/bot{BOT_TOKEN}/sendMessage"
    try:
        proxies = None
        if os.getenv("USE_PROXY", "").lower() == "true":
            proxy_url = f"socks5://{os.getenv('PROXY_HOST', '127.0.0.1')}:{os.getenv('PROXY_PORT', '10808')}"
            proxies = {"http": proxy_url, "https": proxy_url}
        requests.post(url, json={"chat_id": ADMIN_USER_ID, "text": text, "parse_mode": "HTML"}, proxies=proxies, timeout=10)
    except Exception as e: log.error(f"Failed to send Telegram notification: {e}")

# === PDF GENERATION ===

def html_to_pdf_stream(html_content):
    try:
        return HTML(string=html_content).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):
    from fpdf import FPDF
    soup = BeautifulSoup(html_content, 'lxml')
    text = soup.get_text(separator='\n')
    pdf = FPDF()
    pdf.add_page()
    font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"
    if not os.path.exists(font_path): font_path = "/usr/share/fonts/dejavu/DejaVuSans.ttf"
    if os.path.exists(font_path):
        pdf.add_font('DejaVu', '', font_path)
        pdf.set_font('DejaVu', '', 10)
    else: pdf.set_font('helvetica', '', 10)
    for line in [l for l in text.split('\n') if l.strip()]:
        try:
            pdf.multi_cell(0, 6, line)
            pdf.ln(1)
        except:
            pdf.multi_cell(0, 6, line.encode('latin-1', 'replace').decode('latin-1'))
            pdf.ln(1)
    return pdf.output()

def get_image_dims(data):
    import io
    img = Image.open(io.BytesIO(data))
    return img.size[0] * 25.4 / 96, img.size[1] * 25.4 / 96

def convert_attachment_to_pdf(data, filename, content_type):
    from fpdf import FPDF
    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(data)
            scale = min((pdf.w - 20) / img_w, (pdf.h - 20) / img_h, 1.0)
            with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp:
                tmp.write(data); tmp_path = tmp.name
            try: pdf.image(tmp_path, x=10, y=10, w=img_w * scale, h=img_h * scale)
            finally: os.unlink(tmp_path)
            return pdf.output()
        elif ext in ('.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx', '.odt', '.ods'):
            with tempfile.TemporaryDirectory() as tmpdir:
                src = Path(tmpdir) / filename
                src.write_bytes(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 data
        elif ext in ('.txt', '.csv', '.log', '.xml', '.html', '.htm', '.json'):
            try: text = data.decode('utf-8', errors='replace')
            except: text = data.decode('latin-1', errors='replace')
            return html_to_pdf_fpdf(text)
        return None
    except Exception as e:
        log.warning(f"Failed to convert {filename}: {e}")
        return None

def merge_pdfs(pdf_bytes_list):
    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: merger.insert_pdf(fitz.open(stream=pdf_bytes, filetype="pdf"))
        except Exception as e: log.warning(f"Merge error: {e}")
    output = merger.tobytes(); merger.close()
    return output

def build_email_html(msg):
    subject = msg.subject or "(Без темы)"
    from_addr = msg.from_ or "unknown"
    to_addr = ", ".join(msg.to) if msg.to else "unknown"
    date_str = msg.date_str or str(msg.date) if msg.date else "unknown"
    body_html = msg.html or (f"<pre>{msg.text}</pre>" if msg.text else "<p><i>(Нет тела письма)</i></p>")
    return f"<html><head><meta charset='utf-8'><style>body {{ font-family: 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; }} .body {{ line-height: 1.6; }}</style></head><body><div class='header'><h2>{subject}</h2><p><b>От:</b> {from_addr}<br><b>Кому:</b> {to_addr}<br><b>Дата:</b> {date_str}</p></div><div class='body'>{body_html}</div></body></html>"

# === CORE LOGIC ===

def run_export():
    mutt_cfg = parse_mutt_config()
    user, password = mutt_cfg.get('imap_user'), mutt_cfg.get('imap_pass')
    if not user or not password: return

    EXPORT_DIR.mkdir(exist_ok=True); ALL_FOLDER.mkdir(exist_ok=True)
    
    log_id = None
    try:
        conn = sqlite3.connect(DB_PATH); cur = conn.cursor()
        cur.execute("INSERT INTO email_export_logs (status, start_time) VALUES (?, CURRENT_TIMESTAMP)", ('running',))
        log_id = cur.lastrowid; conn.commit(); conn.close()
    except Exception as e: log.error(f"DB Log Error: {e}")

    processed_uids = set()
    if PROGRESS_FILE.exists(): processed_uids = set(PROGRESS_FILE.read_text().splitlines())

    new_emails_info = []

    try:
        with MailBox(IMAP_SERVER).login(user, password) as mailbox:
            all_mail_folder = "[Gmail]/Вся почта"
            for folder in mailbox.folder.list():
                if r'\All' in folder.flags: all_mail_folder = folder.name; break
            mailbox.folder.set(all_mail_folder)
            
            log.info(f"Phase 1: Scanning {all_mail_folder} headers...")
            since_str = SINCE_DATE.strftime("%d-%b-%Y")
            search_criteria = f'(SINCE "{since_str}")'
            
            matching_uids = []
            for msg in mailbox.fetch(search_criteria, headers_only=True, mark_seen=False):
                if str(msg.uid) in processed_uids: continue
                
                from_addr = (msg.from_ or "").lower()
                to_addrs = [t.lower() for t in msg.to] if msg.to else []
                cc_addrs = [c.lower() for c in msg.cc] if msg.cc else []
                subject_lower = (msg.subject or "").lower()
                
                all_targets = [from_addr, subject_lower] + to_addrs + cc_addrs
                found_domain = None
                for d in DOMAINS:
                    if any(d in target for target in all_targets):
                        found_domain = d
                        break
                
                # Check keywords if domain not found
                if not found_domain:
                    for kw in KEYWORDS:
                        if kw.lower() in subject_lower:
                            found_domain = "court_manual"
                            break
                
                if not found_domain: continue

            if not matching_uids:
                log.info("No new matching emails found."); 
            else:
                log.info(f"Phase 2: Downloading {len(matching_uids)} full emails...")
                for uid, domain in matching_uids:
                    try:
                        full_msgs = list(mailbox.fetch(AND(uid=uid), mark_seen=False))
                        if not full_msgs: continue
                        msg = full_msgs[0]
                        
                        pdf_parts = [html_to_pdf_stream(build_email_html(msg))]
                        if msg.attachments:
                            for att in msg.attachments:
                                att_pdf = convert_attachment_to_pdf(att.payload, att.filename, att.content_type)
                                if att_pdf: pdf_parts.append(att_pdf)
                        
                        merged_pdf = merge_pdfs(pdf_parts)
                        if merged_pdf:
                            msg_date = msg.date.strftime("%Y-%m-%d")
                            filename = f"{msg_date}_{domain}_{sanitize_filename(msg.subject)}.pdf"
                            (ALL_FOLDER / filename).write_bytes(merged_pdf)
                            
                            new_emails_info.append({'domain': domain, 'date': msg_date, 'inbound': user.lower() in [t.lower() for t in msg.to] if msg.to else True, 'subject': msg.subject})
                            
                            processed_uids.add(str(uid))
                            with open(PROGRESS_FILE, 'a') as pf: pf.write(f"{uid}\n")
                            log.info(f"Saved: {filename}")
                    except Exception as e: log.error(f"Error fetching UID {uid}: {e}")

    except Exception as e:
        log.error(f"IMAP Error: {e}")
        if log_id:
            conn = sqlite3.connect(DB_PATH); cur = conn.cursor()
            cur.execute("UPDATE email_export_logs SET status = ?, end_time = CURRENT_TIMESTAMP WHERE id = ?", ('error', log_id))
            conn.commit(); conn.close()
        return

    if new_emails_info:
        try:
            conn = sqlite3.connect(DB_PATH); cur = conn.cursor()
            for item in new_emails_info:
                cur.execute("""INSERT INTO email_audit (domain, description, inbound_count, outbound_count, last_email_date, updated_at) VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP) ON CONFLICT(domain) DO UPDATE SET inbound_count = inbound_count + ?, outbound_count = outbound_count + ?, last_email_date = MAX(last_email_date, ?), updated_at = CURRENT_TIMESTAMP""", (item['domain'], item['domain'], 1 if item['inbound'] else 0, 0 if item['inbound'] else 1, item['date'], 1 if item['inbound'] else 0, 0 if item['inbound'] else 1, item['date']))
            conn.commit(); conn.close()
        except Exception as e: log.error(f"DB Update Error: {e}")

        # Merge Master PDF
        pdf_files = sorted(list(ALL_FOLDER.glob("*.pdf")))
        if pdf_files:
            master = fitz.open()
            for pdf in pdf_files:
                try: master.insert_pdf(fitz.open(pdf))
                except Exception as e: log.error(f"Merge error {pdf.name}: {e}")
            master.save(MASTER_PDF); master.close()

        subjects = "\n".join([f"• {e['domain']}: {e['subject'][:50]}..." for e in new_emails_info[:5]])
        if len(new_emails_info) > 5: subjects += f"\n... и еще {len(new_emails_info)-5} писем."
        send_bot_message(f"📧 <b>Email Auditor: Обновление</b>\n\nНайдено новых писем: <b>{len(new_emails_info)}</b>\n\n<b>Последние темы:</b>\n{subjects}\n\nАрхив обновлен.")

    if log_id:
        try:
            conn = sqlite3.connect(DB_PATH); cur = conn.cursor()
            cur.execute("UPDATE email_export_logs SET status = ?, end_time = CURRENT_TIMESTAMP, new_emails_count = ? WHERE id = ?", ('success', len(new_emails_info), log_id))
            conn.commit(); conn.close()
        except: pass

if __name__ == "__main__":
    run_export()