#!/usr/bin/env python3
"""
Gmail Daily Triage — Harvey
Runs on a schedule:
  1. Labels all new unread emails (Action/Read/Receipt/Newsletter/Junk)
  2. Archives everything except Action items
  3. Drafts replies for Action items in Mike's voice → saves to Gmail Drafts
  4. Sends Telegram summary
"""
import imaplib
import email
import smtplib
import json
import re
import urllib.request
import urllib.parse
import time
from email.header import decode_header
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from datetime import datetime
from pathlib import Path

GMAIL_USER  = "mikeziarko@gmail.com"
SECRETS_DIR = Path.home() / ".openclaw/secrets"
TELEGRAM_CHAT_ID = "8792051045"

LABELS = {
    "action":     "Harvey/Action",
    "read":       "Harvey/Read",
    "receipt":    "Harvey/Receipt",
    "newsletter": "Harvey/Newsletter",
    "junk":       "Harvey/Junk",
}

def get_secret(name):
    return (SECRETS_DIR / name).read_text().strip()

def decode_str(s):
    if not s: return ""
    parts = decode_header(s)
    out = ""
    for part, charset in parts:
        if isinstance(part, bytes):
            out += part.decode(charset or "utf-8", errors="replace")
        else:
            out += part
    return out

def get_body(msg):
    text = ""
    if msg.is_multipart():
        for part in msg.walk():
            ct = part.get_content_type()
            cd = str(part.get("Content-Disposition", ""))
            if "attachment" in cd: continue
            if ct == "text/plain":
                try:
                    text += part.get_payload(decode=True).decode(part.get_content_charset() or "utf-8", errors="replace")
                except: pass
            elif ct == "text/html" and not text:
                try:
                    raw = part.get_payload(decode=True).decode(part.get_content_charset() or "utf-8", errors="replace")
                    raw = re.sub(r'<style[^>]*>.*?</style>', ' ', raw, flags=re.DOTALL)
                    raw = re.sub(r'<[^>]+>', ' ', raw)
                    raw = re.sub(r'\s+', ' ', raw).strip()
                    text += raw
                except: pass
    else:
        try:
            text = msg.get_payload(decode=True).decode(msg.get_content_charset() or "utf-8", errors="replace")
        except:
            text = str(msg.get_payload())
    return text[:3000].strip()

def get_imap():
    imap = imaplib.IMAP4_SSL("imap.gmail.com")
    imap.login(GMAIL_USER, get_secret("gmail-app-password.txt"))
    return imap

def claude(prompt, max_tokens=500):
    api_key = get_secret("anthropic-api-key.txt")
    payload = json.dumps({
        "model": "claude-haiku-4-5",
        "max_tokens": max_tokens,
        "messages": [{"role": "user", "content": prompt}]
    }).encode()
    req = urllib.request.Request(
        "https://api.anthropic.com/v1/messages",
        data=payload,
        headers={"x-api-key": api_key, "anthropic-version": "2023-06-01", "content-type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=20) as r:
        return json.loads(r.read())["content"][0]["text"].strip()

def categorize_batch(emails_meta):
    lines = "\n".join(
        f'{i+1}. From: {e["sender"][:60]} | Subject: {e["subject"][:80]} | Preview: {e["snippet"][:100]}'
        for i, e in enumerate(emails_meta)
    )
    prompt = f"""Categorize each email for Mike Ziarko, owner of No More Chores (Toronto cleaning company).

Categories:
- action: needs a reply or decision (customer issues, team questions, vendor problems, anything requiring response)
- read: informational, worth reading, no reply needed (alerts, business updates, articles)
- receipt: payment confirmations, invoices, order confirmations, bills
- newsletter: marketing emails, newsletters, promotional content
- junk: spam, cold outreach, sweepstakes, irrelevant promotions

Reply ONLY with JSON mapping line number to category. Example: {{"1":"action","2":"receipt"}}

Emails:
{lines}

JSON:"""
    try:
        result = claude(prompt, max_tokens=600)
        m = re.search(r'\{{[^{{}}]+\}}', result, re.DOTALL)
        if m:
            return json.loads(m.group())
    except Exception as e:
        print(f"  Categorize error: {e}")
    return {}

def draft_reply(sender, subject, body):
    prompt = f"""Draft a concise, warm reply on behalf of Mike Ziarko, owner of No More Chores (Toronto residential cleaning company).

Mike's voice:
- Direct and friendly, not corporate
- Gets to the point quickly
- Signs off as Mike
- Never uses em dashes
- Warm but not sycophantic

From: {sender}
Subject: {subject}

Email content:
{body}

Write a complete reply Mike can send with minimal editing. Just the reply body, no subject line:"""
    try:
        return claude(prompt, max_tokens=400)
    except:
        return ""

def save_draft(to_addr, subject, body, in_reply_to=None):
    msg = MIMEMultipart()
    msg["From"]    = GMAIL_USER
    msg["To"]      = to_addr
    msg["Subject"] = f"Re: {subject}" if not subject.lower().startswith("re:") else subject
    msg["Date"]    = email.utils.formatdate(localtime=True)
    if in_reply_to:
        msg["In-Reply-To"] = in_reply_to
        msg["References"]  = in_reply_to
    msg.attach(MIMEText(body, "plain"))

    imap = get_imap()
    imap.append("[Gmail]/Drafts", "\\Draft", imaplib.Time2Internaldate(time.time()), msg.as_bytes())
    imap.logout()

def send_telegram(message):
    config_path = Path.home() / ".openclaw/openclaw.json"
    try:
        config = json.loads(config_path.read_text())
        token = None
        for acct in config.get("accounts", []):
            if acct.get("provider") == "telegram":
                token = acct.get("botToken") or acct.get("token")
                if token: break
        if not token:
            print("No Telegram token found")
            return
    except Exception as e:
        print(f"Config error: {e}")
        return

    payload = json.dumps({
        "chat_id": TELEGRAM_CHAT_ID,
        "text": message,
        "parse_mode": "HTML"
    }).encode()
    req = urllib.request.Request(
        f"https://api.telegram.org/bot{token}/sendMessage",
        data=payload,
        headers={"content-type": "application/json"}
    )
    try:
        urllib.request.urlopen(req, timeout=10)
    except Exception as e:
        print(f"Telegram error: {e}")

def main():
    print(f"Daily triage starting at {datetime.now().strftime('%H:%M')}...")
    imap = get_imap()
    imap.select("INBOX")

    # Ensure labels exist
    for label in LABELS.values():
        try: imap.create(label)
        except: pass

    # Get unread emails
    _, msgs = imap.uid('SEARCH', None, 'UNSEEN')
    all_uids = msgs[0].split()
    if not all_uids:
        print("Inbox is clean. Nothing to triage.")
        imap.logout()
        return

    print(f"Found {len(all_uids)} unread emails...")

    # Fetch headers in batches
    emails = []
    BATCH = 25
    for i in range(0, len(all_uids), BATCH):
        batch = all_uids[i:i+BATCH]
        uid_str = b','.join(batch)
        _, data = imap.uid('FETCH', uid_str, '(BODY.PEEK[HEADER.FIELDS (FROM SUBJECT MESSAGE-ID LIST-UNSUBSCRIBE LIST-ID)])')
        for j in range(0, len(data), 2):
            try:
                uid_line = data[j][0].decode()
                uid = uid_line.split('UID ')[1].split(' ')[0] if 'UID ' in uid_line else None
                if not uid: continue
                msg = email.message_from_bytes(data[j][1])
                sender  = decode_str(msg.get("From", ""))
                subject = decode_str(msg.get("Subject", "(no subject)"))
                unsub   = msg.get("List-Unsubscribe", "")
                msg_id  = msg.get("Message-ID", "")
                emails.append({
                    "uid":     uid,
                    "sender":  sender,
                    "subject": subject,
                    "snippet": "",
                    "unsub":   unsub,
                    "msg_id":  msg_id
                })
            except: pass

    imap.logout()

    # Categorize
    print("Categorizing...")
    results = {}
    for i in range(0, len(emails), BATCH):
        batch = emails[i:i+BATCH]
        mapping = categorize_batch(batch)
        for j, e in enumerate(batch):
            cat = mapping.get(str(j+1), "read")
            if cat not in LABELS: cat = "read"
            if e["unsub"] and cat not in ("action", "receipt"):
                cat = "newsletter"
            results[e["uid"]] = cat

    # Apply labels
    print("Applying labels...")
    imap = get_imap()
    imap.select("INBOX")

    by_cat = {c: [] for c in LABELS}
    for e in emails:
        cat = results.get(e["uid"], "read")
        by_cat[cat].append(e)

    for cat, label in LABELS.items():
        uids = [e["uid"].encode() for e in by_cat[cat]]
        if uids:
            uid_str = b','.join(uids)
            imap.uid('STORE', uid_str, '+X-GM-LABELS', f'"{label}"')

    # Archive non-action from inbox
    non_action = [e["uid"].encode() for e in emails if results.get(e["uid"]) != "action"]
    if non_action:
        uid_str = b','.join(non_action)
        imap.uid('COPY', uid_str, '"[Gmail]/All Mail"')
        imap.uid('STORE', uid_str, '+FLAGS', '\\Deleted')
        imap.expunge()

    imap.logout()

    action_emails = by_cat["action"]
    print(f"Action items: {len(action_emails)}, drafting replies...")

    # Draft replies for action items
    drafted = []
    if action_emails:
        imap = get_imap()
        imap.select('"[Gmail]/All Mail"')
        for e in action_emails:
            try:
                _, data = imap.uid('FETCH', e["uid"].encode(), '(RFC822)')
                full_msg = email.message_from_bytes(data[0][1])
                body = get_body(full_msg)
                reply = draft_reply(e["sender"], e["subject"], body)
                if reply:
                    sender_addr = re.search(r'<([^>]+)>', e["sender"])
                    to = sender_addr.group(1) if sender_addr else e["sender"]
                    save_draft(to, e["subject"], reply, e["msg_id"])
                    drafted.append(e["subject"][:60])
                    print(f"  Drafted reply: {e['subject'][:50]}")
            except Exception as ex:
                print(f"  Draft error: {ex}")
        imap.logout()

    # Telegram summary
    counts = {c: len(by_cat[c]) for c in LABELS}
    total = len(emails)

    lines = [f"<b>Daily Triage Complete</b> -- {datetime.now().strftime('%-I:%M %p')}"]
    lines.append(f"{total} emails processed\n")
    lines.append(f"🔴 Action: {counts['action']} ({len(drafted)} drafts ready in Gmail)")
    lines.append(f"📖 Read: {counts['read']}")
    lines.append(f"💰 Receipts: {counts['receipt']}")
    lines.append(f"📰 Newsletters: {counts['newsletter']}")
    lines.append(f"🗑 Junk: {counts['junk']}")

    if drafted:
        lines.append("\n<b>Drafts ready to review:</b>")
        for d in drafted[:5]:
            lines.append(f"  • {d}")

    send_telegram("\n".join(lines))
    print("Done. Telegram summary sent.")

if __name__ == "__main__":
    main()