import imaplib
import os
import logging
import sys
from dotenv import load_dotenv
from datetime import datetime, timezone, timedelta
from email import message_from_bytes

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)

load_dotenv()

# Configuration
source_host = os.getenv("SOURCE_HOST")
if not source_host:
    raise ValueError("SOURCE_HOST not set")
source_port = int(os.getenv("SOURCE_PORT", "993"))
source_user = os.getenv("SOURCE_USER")
if not source_user:
    raise ValueError("SOURCE_USER not set")
source_pass = os.getenv("SOURCE_PASS")
if not source_pass:
    raise ValueError("SOURCE_PASS not set")
dest_host = os.getenv("DEST_HOST")
if not dest_host:
    raise ValueError("DEST_HOST not set")
dest_port = int(os.getenv("DEST_PORT", "993"))
dest_user = os.getenv("DEST_USER")
if not dest_user:
    raise ValueError("DEST_USER not set")
dest_pass = os.getenv("DEST_PASS")
if not dest_pass:
    raise ValueError("DEST_PASS not set")
folders_config = os.getenv("FOLDERS", "INBOX")
processed_msg_file = "processed_message_ids.txt"
last_run_file = "last_run.txt"
dry_run = os.getenv("DRY_RUN", "false").lower() == "true"


def main():
    # Last run
    if os.path.exists(last_run_file):
        with open(last_run_file) as f:
            last_run_str = f.read().strip()
        last_run = datetime.fromisoformat(last_run_str)
    else:
        last_run = datetime.now(timezone.utc) - timedelta(hours=1)
    since_date = last_run.strftime("%d-%b-%Y")
    logging.info(f"Using SINCE {since_date}")

    # Load processed Message-IDs
    processed_msg = set()
    if os.path.exists(processed_msg_file):
        with open(processed_msg_file, "r") as f:
            processed_msg = set(line.strip() for line in f if line.strip())

    try:
        # Connect to source IMAP for listing
        source_imap_list = imaplib.IMAP4_SSL(source_host, source_port)
        source_imap_list.login(source_user, source_pass)

        # Determine folders to sync
        if folders_config == "all":
            # List all folders
            typ, folder_list = source_imap_list.list()
            if typ != "OK":
                raise Exception("Failed to list folders")
            folders = []
            for f in folder_list:
                if f:
                    # Parse folder name, e.g., '(\\HasNoChildren) "/" "INBOX"'
                    parts = f.decode().split(' "/" ')
                    if len(parts) > 1:
                        folder_name = parts[1].strip('"')
                        # Skip system folders if desired (optional)
                        if folder_name not in ["Trash", "Junk", "Spam"]:
                            folders.append(folder_name)
        else:
            folders = [f.strip() for f in folders_config.split(",")]

        source_imap_list.logout()

        logging.info(f"Syncing folders: {folders}")

        # Connect to destination IMAP
        dest_imap = imaplib.IMAP4_SSL(dest_host, dest_port)
        dest_imap.login(dest_user, dest_pass)

        total_forwarded = 0

        for folder in folders:
            try:
                # Connect to source IMAP per folder
                source_imap = imaplib.IMAP4_SSL(source_host, source_port)
                source_imap.login(source_user, source_pass)
                source_imap.select(folder)

                # Search for emails since last run
                typ, data = source_imap.search(None, f'SINCE "{since_date}"')
                if typ != "OK":
                    logging.error(f"Search failed for {folder}: {typ}")
                    source_imap.logout()
                    continue

                uids = data[0].split()
                logging.info(f"Found {len(uids)} emails in {folder} since {since_date}")

                # Ensure dest folder exists
                try:
                    dest_imap.select(folder)
                except:
                    logging.info(f"Creating folder {folder} in dest")
                    dest_imap.create(folder)
                    dest_imap.select(folder)

                forwarded = 0
                for uid in uids:
                    uid_str = uid.decode()
                    # Fetch raw message
                    typ, msg_data = source_imap.fetch(uid, "(RFC822)")
                    if typ != "OK":
                        logging.error(f"Failed to fetch UID {uid_str} in {folder}")
                        continue
                    raw_msg = msg_data[0][1]

                    msg = message_from_bytes(raw_msg)
                    msg_id = msg["Message-ID"]
                    if not msg_id:
                        logging.warning(
                            f"No Message-ID for UID {uid_str} in {folder}, skipping"
                        )
                        continue
                    if msg_id in processed_msg:
                        logging.info(f"Skipping already processed Message-ID {msg_id}")
                        continue

                    # Check dest for duplicate
                    typ, dest_data = dest_imap.search(
                        None, f'HEADER Message-ID "{msg_id}"'
                    )
                    if typ == "OK" and dest_data[0]:
                        logging.info(f"Skipping duplicate Message-ID {msg_id} in dest")
                        continue

                    if not dry_run:
                        # Append to destination
                        dest_imap.append(folder, "", None, raw_msg)

                        # Mark as processed
                        with open(processed_msg_file, "a") as f:
                            f.write(msg_id + "\n")
                        processed_msg.add(msg_id)
                        forwarded += 1
                        logging.info(
                            f"Forwarded email Message-ID {msg_id} from {folder}"
                        )
                    else:
                        logging.info(
                            f"Dry-run: Would forward Message-ID {msg_id} from {folder}"
                        )

                source_imap.logout()
                total_forwarded += forwarded
                logging.info(
                    f"Completed syncing {folder}: forwarded {forwarded} emails"
                )

            except Exception as e:
                logging.error(f"Error syncing folder {folder}: {e}")

        dest_imap.logout()

        if not dry_run:
            # Update last run
            with open(last_run_file, "w") as f:
                f.write(datetime.now(timezone.utc).isoformat())

        logging.info(
            f"Email forwarding completed. Total forwarded {total_forwarded} new emails."
        )

    except Exception as e:
        logging.error(f"Error during email forwarding: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()