import imaplib
import email
from email.header import decode_header
import os
from pathlib import Path
import paramiko

# =========================
# EMAIL CONFIG
# =========================
IMAP_SERVER = "outlook.office365.com"
EMAIL_ACCOUNT = "mailbox@company.com"
PASSWORD = "your_password"

MAIL_FOLDERS = ["Inbox", "Orders"]

# =========================
# SFTP CONFIG
# =========================
SFTP_HOST = "sftp.company.com"
SFTP_PORT = 22
SFTP_USERNAME = "sftpuser"
SSH_KEY_PATH = "/home/youruser/.ssh/id_rsa"
SSH_KEY_PASSPHRASE = None

# =========================
# RULES
# =========================
RULES = [
    {
        "from_contains": "supplier1@company.com",
        "to_contains": "orders@yourcompany.com",
        "remote_path": "/incoming/orders"
    },
    {
        "from_contains": "reports@vendor.com",
        "to_contains": "reports@yourcompany.com",
        "remote_path": "/incoming/reports"
    }
]

TEMP_DOWNLOAD = "/tmp/mail_attachments"

# =========================
# HELPERS
# =========================
def ensure_folder(path):
    Path(path).mkdir(parents=True, exist_ok=True)

def clean_subject(subject):
    if not subject:
        return "No Subject"
    parts = decode_header(subject)
    decoded_subject = ""
    for part, encoding in parts:
        if isinstance(part, bytes):
            decoded_subject += part.decode(encoding or "utf-8", errors="ignore")
        else:
            decoded_subject += part
    return decoded_subject

def decode_filename(filename):
    if not filename:
        return "unknown"
    parts = decode_header(filename)
    decoded_name = ""
    for part, encoding in parts:
        if isinstance(part, bytes):
            decoded_name += part.decode(encoding or "utf-8", errors="ignore")
        else:
            decoded_name += part
    return decoded_name

# =========================
# SFTP CONNECTION
# =========================
def create_sftp_connection():
    try:
        key = paramiko.RSAKey.from_private_key_file(
            SSH_KEY_PATH,
            password=SSH_KEY_PASSPHRASE
        )
        transport = paramiko.Transport((SFTP_HOST, SFTP_PORT))
        transport.connect(username=SFTP_USERNAME, pkey=key)
        sftp = paramiko.SFTPClient.from_transport(transport)
        return sftp, transport
    except Exception as e:
        print(f"SFTP connection failed: {e}")
        return None, None

# =========================
# UPLOAD FILE
# =========================
def upload_to_sftp(local_file, remote_folder):
    sftp, transport = create_sftp_connection()
    if not sftp:
        print("Skipping upload due to SFTP connection error")
        return
    try:
        filename = os.path.basename(local_file)
        remote_path = f"{remote_folder}/{filename}"
        print(f"Uploading {filename} -> {remote_path}")
        sftp.put(local_file, remote_path)
        print("Upload successful")
    except Exception as e:
        print(f"Failed to upload {filename}: {e}")
    finally:
        sftp.close()
        transport.close()

# =========================
# PROCESS EMAIL
# =========================
def process_email(msg):
    from_addr = msg.get("From", "").lower()
    to_addr = msg.get("To", "").lower()
    subject = clean_subject(msg.get("Subject"))

    print("\n====================")
    print(f"FROM: {from_addr}")
    print(f"TO: {to_addr}")
    print(f"SUBJECT: {subject}")

    matching_rule = None
    for rule in RULES:
        if rule["from_contains"].lower() in from_addr and rule["to_contains"].lower() in to_addr:
            matching_rule = rule
            break

    if not matching_rule:
        print("No matching rule")
        return

    remote_folder = matching_rule["remote_path"]

    for part in msg.walk():
        content_disposition = str(part.get("Content-Disposition"))
        if "attachment" not in content_disposition.lower():
            continue

        filename = decode_filename(part.get_filename())
        ensure_folder(TEMP_DOWNLOAD)
        local_path = os.path.join(TEMP_DOWNLOAD, filename)
        with open(local_path, "wb") as f:
            f.write(part.get_payload(decode=True))
        print(f"Saved attachment locally: {local_path}")

        upload_to_sftp(local_path, remote_folder)

# =========================
# MAIN
# =========================
def main():
    try:
        mail = imaplib.IMAP4_SSL(IMAP_SERVER)
        mail.login(EMAIL_ACCOUNT, PASSWORD)
        print("Connected to mailbox")
    except Exception as e:
        print(f"Failed to connect/login: {e}")
        return

    for folder in MAIL_FOLDERS:
        print(f"\nChecking folder: {folder}")
        status, _ = mail.select(folder)
        if status.decode() != "OK":
            print(f"Cannot access folder: {folder}")
            continue

        status, messages = mail.search(None, "UNSEEN")
        if status.decode() != "OK":
            continue

        email_ids = messages[0].split()
        print(f"Found {len(email_ids)} unread emails")

        for email_id in email_ids:
            status, msg_data = mail.fetch(email_id, "(RFC822)")
            if status.decode() != "OK":
                continue
            raw_email = msg_data[0][1]
            msg = email.message_from_bytes(raw_email)
            process_email(msg)
            mail.store(email_id, '+FLAGS', '\\Seen')

    mail.logout()
    print("\nFinished")

if __name__ == "__main__":
    main()