#!/usr/bin/env python3
import email
import imaplib
import json
import logging
import os
import re
from datetime import datetime, timedelta, timezone

import redis
from celery import shared_task

from app.config import settings
from app.tasks.convert_to_pdf import convert_to_pdf  # new conversion task
from app.tasks.process_document import process_document  # Updated import
from app.utils.allowed_types import (
    ALL_CATEGORIES,
    DEFAULT_CATEGORIES,
    get_allowed_types_for_categories,
)

# Database session for per-user IMAP accounts (imported lazily to avoid circular imports)
_db_session_factory = None

# Maximum length to store as last_error to prevent DB bloat
_MAX_ERROR_LENGTH = 500


def _get_db_session():
    """Return a new SQLAlchemy session (lazy import to avoid startup issues)."""
    global _db_session_factory  # noqa: PLW0603
    if _db_session_factory is None:
        from app.database import SessionLocal

        _db_session_factory = SessionLocal
    return _db_session_factory()


logger = logging.getLogger(__name__)

# Initialize Redis connection using Celery's Redis settings
redis_client = redis.StrictRedis.from_url(settings.redis_url, decode_responses=True)


def _decrypt_imap_password(password: str | None) -> str | None:
    """Decrypt an IMAP account password stored in the database.

    Passwords are stored encrypted (Fernet, ``enc:`` prefix) for new records;
    legacy plaintext records are returned unchanged so existing accounts
    continue to work until they are next updated via the API.
    """
    from app.utils.encryption import decrypt_value

    return decrypt_value(password)


def _resolve_categories_for_profile(profile_id: int | None) -> list[str]:
    """Return the list of allowed categories for a profile ID.

    Loads the profile from the database.  If ``profile_id`` is ``None`` or the
    profile is not found, falls back to the global ``settings.imap_attachment_filter``
    string (``'documents_only'`` → default categories; ``'all'`` → all categories).
    """
    if profile_id is not None:
        try:
            from app.models import ImapIngestionProfile

            db = _get_db_session()
            try:
                profile = db.query(ImapIngestionProfile).filter(ImapIngestionProfile.id == profile_id).first()
                if profile:
                    return json.loads(profile.allowed_categories)
            finally:
                db.close()
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "Could not load IMAP ingestion profile %d (%s: %s) — using global default",
                profile_id,
                type(exc).__name__,
                exc,
            )

    # Fall back to global setting
    if settings.imap_attachment_filter == "all":
        return ALL_CATEGORIES
    return DEFAULT_CATEGORIES


LOCK_KEY = "imap_lock"  # Unique key for locking
LOCK_EXPIRE = 300  # Lock expires in 5 minutes

# Local cache file for tracking processed emails
CACHE_FILE = os.path.join(settings.workdir, "processed_mails.json")


def acquire_lock():
    """Attempt to acquire a Redis-based lock. If acquired, set an expiration."""
    lock_acquired = redis_client.setnx(LOCK_KEY, "locked")
    if lock_acquired:
        redis_client.expire(LOCK_KEY, LOCK_EXPIRE)
        logger.info("Lock acquired for IMAP processing.")
        return True
    logger.warning("Lock already held. Skipping this cycle.")
    return False


def release_lock():
    """Release the lock by deleting the Redis key."""
    redis_client.delete(LOCK_KEY)
    logger.info("Lock released.")


def load_processed_emails():
    """Load the list of already processed emails from a local JSON file."""
    if os.path.exists(CACHE_FILE):
        try:
            with open(CACHE_FILE, "r") as f:
                processed_emails = json.load(f)
                processed_emails = cleanup_old_entries(processed_emails)
                return processed_emails
        except json.JSONDecodeError:
            logger.warning("Failed to decode JSON, resetting processed emails cache.")
            return {}
    return {}


def save_processed_emails(processed_emails):
    """Save the processed email IDs to a local JSON file."""
    with open(CACHE_FILE, "w") as f:
        json.dump(processed_emails, f, indent=4)


def cleanup_old_entries(processed_emails):
    """Remove entries older than 7 days from the cache to avoid infinite growth."""
    seven_days_ago = datetime.now(timezone.utc) - timedelta(days=7)
    valid_emails = {}
    for msg_id, date_str in processed_emails.items():
        naive_dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S")
        aware_dt = naive_dt.replace(tzinfo=timezone.utc)
        if aware_dt > seven_days_ago:
            valid_emails[msg_id] = date_str
    return valid_emails


@shared_task
def pull_all_inboxes():
    """
    Periodic Celery task that checks all configured IMAP mailboxes
    and fetches attachments from new emails.
    Ensures only one instance runs at a time using Redis-based locking.

    Processes:
    1. System-level mailboxes configured via environment variables (IMAP1, IMAP2).
    2. Per-user IMAP accounts stored in the ``user_imap_accounts`` database table.
    """
    if not acquire_lock():
        logger.info("Skipping execution: Another instance is running.")
        return

    try:
        logger.info("Starting pull_all_inboxes")

        # Mailbox #1 (non-Gmail)
        check_and_pull_mailbox(
            mailbox_key="imap1",
            host=settings.imap1_host,
            port=settings.imap1_port,
            username=settings.imap1_username,
            password=settings.imap1_password,
            use_ssl=settings.imap1_ssl,
            delete_after_process=settings.imap1_delete_after_process,
        )

        # Mailbox #2 (Gmail)
        check_and_pull_mailbox(
            mailbox_key="imap2",
            host=settings.imap2_host,
            port=settings.imap2_port,
            username=settings.imap2_username,
            password=settings.imap2_password,
            use_ssl=settings.imap2_ssl,
            delete_after_process=settings.imap2_delete_after_process,
        )

        # Per-user IMAP accounts from the database
        _pull_user_imap_accounts()

        # Per-user IMAP integrations from the UserIntegration model
        _pull_user_integration_imap()

        logger.info("Finished pull_all_inboxes")

    finally:
        release_lock()


def _pull_user_imap_accounts() -> None:
    """Iterate over all active per-user IMAP accounts and pull their inboxes."""
    try:
        from app.models import UserImapAccount

        db = _get_db_session()
        try:
            accounts = db.query(UserImapAccount).filter(UserImapAccount.is_active.is_(True)).all()
            logger.info("Processing %d per-user IMAP account(s)", len(accounts))
            for acct in accounts:
                # Use a descriptive identifier for logging and processed-email cache keys
                account_identifier = f"user_{acct.owner_id}_{acct.id}"
                try:
                    pull_inbox(
                        mailbox_key=account_identifier,
                        host=acct.host,
                        port=acct.port,
                        username=acct.username,
                        password=_decrypt_imap_password(acct.password),
                        use_ssl=acct.use_ssl,
                        delete_after_process=acct.delete_after_process,
                        owner_id=acct.owner_id,
                        allowed_categories=_resolve_categories_for_profile(acct.profile_id),
                    )
                    # Record successful poll
                    acct.last_checked_at = datetime.now(timezone.utc)
                    acct.last_error = None
                    db.commit()
                except Exception as exc:  # noqa: BLE001
                    error_msg = str(exc)[:_MAX_ERROR_LENGTH]
                    logger.error(
                        "Error pulling user IMAP account %d (%s@%s): %s",
                        acct.id,
                        acct.username,
                        acct.host,
                        error_msg,
                    )
                    try:
                        acct.last_checked_at = datetime.now(timezone.utc)
                        acct.last_error = error_msg
                        db.commit()
                    except Exception:  # noqa: BLE001
                        db.rollback()
        finally:
            db.close()
    except Exception as exc:  # noqa: BLE001
        logger.error("Failed to process per-user IMAP accounts: %s", exc)


def _pull_user_integration_imap() -> None:
    """Iterate over all active IMAP UserIntegrations and pull their inboxes.

    This polls the ``user_integrations`` table for records with
    ``integration_type='IMAP'``, ``direction='SOURCE'``, and ``is_active=True``.
    Each integration's config/credentials are decoded and passed to
    :func:`pull_inbox` with the owning user's ``owner_id`` so that ingested
    documents are correctly attributed.

    Individual connection failures are caught and recorded on the integration
    without crashing the polling loop.
    """
    try:
        import json as _json

        from app.models import IntegrationDirection, IntegrationType, UserIntegration
        from app.utils.encryption import decrypt_value

        db = _get_db_session()
        try:
            integrations = (
                db.query(UserIntegration)
                .filter(
                    UserIntegration.integration_type == IntegrationType.IMAP,
                    UserIntegration.direction == IntegrationDirection.SOURCE,
                    UserIntegration.is_active.is_(True),
                )
                .all()
            )
            logger.info("Processing %d IMAP UserIntegration(s)", len(integrations))
            for integ in integrations:
                account_identifier = f"integration_{integ.owner_id}_{integ.id}"
                try:
                    cfg = _json.loads(integ.config) if integ.config else {}
                    raw_creds = decrypt_value(integ.credentials) if integ.credentials else None
                    creds = _json.loads(raw_creds) if raw_creds else {}

                    host = cfg.get("host")
                    port = int(cfg.get("port", 993))
                    username = cfg.get("username")
                    password = creds.get("password")
                    use_ssl = cfg.get("use_ssl", True)
                    delete_after = cfg.get("delete_after_process", False)
                    gmail_labels = cfg.get("gmail_apply_labels", True)
                    # Integrations can store a profile_id in config; fall back to global default
                    profile_id = cfg.get("profile_id")
                    allowed_categories = _resolve_categories_for_profile(profile_id)

                    if not (host and username and password):
                        logger.warning(
                            "IMAP integration %d (owner %s) has incomplete config — skipping.",
                            integ.id,
                            integ.owner_id,
                        )
                        continue

                    pull_inbox(
                        mailbox_key=account_identifier,
                        host=host,
                        port=port,
                        username=username,
                        password=password,
                        use_ssl=use_ssl,
                        delete_after_process=delete_after,
                        owner_id=integ.owner_id,
                        gmail_apply_labels=gmail_labels,
                        allowed_categories=allowed_categories,
                    )
                    integ.last_used_at = datetime.now(timezone.utc)
                    integ.last_error = None
                    db.commit()
                except Exception as exc:  # noqa: BLE001
                    error_msg = str(exc)[:_MAX_ERROR_LENGTH]
                    logger.error(
                        "Error pulling IMAP integration %d (owner %s): %s",
                        integ.id,
                        integ.owner_id,
                        error_msg,
                    )
                    try:
                        integ.last_used_at = datetime.now(timezone.utc)
                        integ.last_error = error_msg
                        db.commit()
                    except Exception:  # noqa: BLE001
                        db.rollback()
        finally:
            db.close()
    except Exception as exc:  # noqa: BLE001
        logger.error("Failed to process IMAP UserIntegrations: %s", exc)


def check_and_pull_mailbox(
    mailbox_key: str,
    host: str | None,
    port: int | None,
    username: str | None,
    password: str | None,
    use_ssl: bool,
    delete_after_process: bool,
):
    """Validates config and invokes pulling from the mailbox if valid."""
    if not (host and port and username and password):
        logger.warning(f"Mailbox {mailbox_key} is missing config, skipping.")
        return

    logger.info(f"Checking mailbox: {mailbox_key}")
    pull_inbox(
        mailbox_key=mailbox_key,
        host=host,
        port=port,
        username=username,
        password=password,
        use_ssl=use_ssl,
        delete_after_process=delete_after_process,
    )


def pull_inbox(
    mailbox_key,
    host,
    port,
    username,
    password,
    use_ssl,
    delete_after_process,
    owner_id=None,
    gmail_apply_labels=True,
    allowed_categories=None,
):
    """
    Connects to the IMAP inbox, fetches new unread emails from the last 3 days,
    and processes attachments while preserving the original unread status.

    For Gmail:
      - Attempts to select the localized All Mail folder.
      - Runs an X-GM-RAW query: "in:anywhere in:unread newer_than:3d has:attachment".

    For non-Gmail mailboxes, it falls back to selecting the INBOX with a SINCE/UNSEEN filter.

    Args:
        owner_id: Optional user identifier. When provided, ingested documents are
                  attributed to this user via ``process_document`` / ``convert_to_pdf``.
        gmail_apply_labels: Whether to apply Gmail-specific labels and stars to
                  processed emails. Only relevant for Gmail hosts. Defaults to True.
        allowed_categories: List of file-type category keys to ingest (e.g.
                  ``["pdf", "office", "images"]``).  ``None`` falls back to the
                  global ``settings.imap_attachment_filter`` mapping.
    """
    if allowed_categories is None:
        allowed_categories = _resolve_categories_for_profile(None)

    effective_mime_types, effective_extensions = get_allowed_types_for_categories(allowed_categories)
    logger.info(
        "Connecting to %s at %s:%s (SSL=%s) — categories: %s",
        mailbox_key,
        host,
        port,
        use_ssl,
        allowed_categories,
    )
    processed_emails = load_processed_emails()

    try:
        mail = imaplib.IMAP4_SSL(host, port) if use_ssl else imaplib.IMAP4(host, port)
        mail.login(username, password)

        is_gmail_host = "gmail" in host.lower()
        if is_gmail_host:
            # For Gmail, try to select the localized All Mail folder.
            all_mail_folder = find_all_mail_folder(mail)
            if all_mail_folder:
                logger.info("Using Gmail All Mail folder: %s", all_mail_folder)
                mail.select(f'"{all_mail_folder}"')
            else:
                logger.warning("Gmail All Mail folder not found, falling back to INBOX.")
                mail.select("INBOX")
            # Use the X-GM-RAW query for Gmail.
            raw_query = "in:anywhere in:unread newer_than:3d has:attachment"
            status, search_data = mail.search(None, "X-GM-RAW", f'"{raw_query}"')
        else:
            # For non-Gmail, select INBOX and use SINCE/UNSEEN query.
            mail.select("INBOX")
            since_date = (datetime.now(timezone.utc) - timedelta(days=3)).strftime("%d-%b-%Y")
            status, search_data = mail.search(None, f"(SINCE {since_date} UNSEEN)")

        if status != "OK":
            logger.warning("Search failed on mailbox %s. Status=%s", mailbox_key, status)
            mail.close()
            mail.logout()
            return

        msg_numbers = search_data[0].split()
        logger.info("Found %d unread emails in %s.", len(msg_numbers), mailbox_key)

        for num in msg_numbers:
            status, msg_data = mail.fetch(num, "(RFC822)")
            if status != "OK":
                logger.warning("Failed to fetch message %s in %s. Status=%s", num, mailbox_key, status)
                continue

            raw_email = msg_data[0][1]
            email_message = email.message_from_bytes(raw_email)
            msg_id = email_message.get("Message-ID")

            if not msg_id:
                logger.warning("Skipping email without Message-ID in %s", mailbox_key)
                continue

            if msg_id in processed_emails:
                logger.info("Skipping already processed email %s in %s", msg_id, mailbox_key)
                continue

            # For Gmail, check if the email already has the "Ingested" label.
            if is_gmail_host and gmail_apply_labels:
                if email_already_has_label(mail, num, "Ingested"):
                    logger.info("Skipping email %s in %s, already labeled 'Ingested'.", msg_id, mailbox_key)
                    continue

            # Process attachments using the resolved mime types / extensions.
            fetch_attachments_and_enqueue(
                email_message,
                owner_id=owner_id,
                effective_mime_types=effective_mime_types,
                effective_extensions=effective_extensions,
            )

            if settings.imap_readonly_mode:
                logger.info("Readonly mode: skipping mailbox modifications for %s in %s", msg_id, mailbox_key)
            else:
                if is_gmail_host and gmail_apply_labels:
                    mark_as_processed_with_star(mail, num)
                    mark_as_processed_with_label(mail, num, label="Ingested")

                if delete_after_process:
                    logger.info("Deleting message %s from %s", num.decode(), mailbox_key)
                    mail.store(num, "+FLAGS", "\\Deleted")
                else:
                    mail.store(num, "-FLAGS", "\\Seen")

            processed_emails[msg_id] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S")
            save_processed_emails(processed_emails)

        if not settings.imap_readonly_mode and delete_after_process:
            mail.expunge()

        mail.close()
        mail.logout()
        logger.info("Finished processing mailbox %s", mailbox_key)

    except Exception as e:
        logger.exception("Error pulling mailbox %s: %s", mailbox_key, e)


def fetch_attachments_and_enqueue(
    email_message,
    owner_id: str | None = None,
    effective_mime_types: frozenset[str] | None = None,
    effective_extensions: frozenset[str] | None = None,
):
    """
    Extracts attachments from the email and processes only allowed file types.

    The caller is responsible for computing ``effective_mime_types`` and
    ``effective_extensions`` from the relevant :class:`ImapIngestionProfile` (or
    the global default) via :func:`app.utils.allowed_types.get_allowed_types_for_categories`
    before calling this function.  ``pull_inbox`` does this automatically.

    If either set is ``None`` the function falls back to the default category list
    so the function still works correctly when called directly in tests or from
    other contexts.

    If the attachment is a PDF (by extension or MIME type), it is enqueued for upload;
    any other allowed file is enqueued for conversion to PDF.

    Args:
        email_message: The parsed email message to extract attachments from.
        owner_id: Optional user identifier forwarded to ``process_document`` /
                  ``convert_to_pdf`` for multi-tenant attribution.
        effective_mime_types: Pre-computed frozenset of allowed MIME type strings.
        effective_extensions: Pre-computed frozenset of allowed file extension strings.

    Returns True if at least one allowed attachment was processed.
    """
    if effective_mime_types is None or effective_extensions is None:
        effective_mime_types, effective_extensions = get_allowed_types_for_categories(DEFAULT_CATEGORIES)

    has_attachment = False
    for part in email_message.walk():
        if part.get_content_maintype() == "multipart":
            continue

        filename = part.get_filename()
        if not filename:
            continue

        # Check if it's a PDF file by extension, regardless of MIME type
        is_pdf_by_extension = filename.lower().endswith(".pdf")

        mime_type = part.get_content_type()
        file_ext = os.path.splitext(filename)[1].lower()

        # Accept file if it has an allowed MIME type, an allowed extension, OR is a PDF by extension
        if mime_type not in effective_mime_types and file_ext not in effective_extensions and not is_pdf_by_extension:
            logger.info(
                "Skipping attachment %s (MIME: %s, ext: %s) — not in effective allowed set",
                filename,
                mime_type,
                file_ext,
            )
            continue

        file_path = os.path.join(settings.workdir, filename)
        with open(file_path, "wb") as f:
            f.write(part.get_payload(decode=True))

        # If it's a PDF by MIME type or extension, process it directly
        if mime_type == "application/pdf" or is_pdf_by_extension:
            process_document.delay(file_path, owner_id=owner_id)
            logger.info("Enqueued PDF for upload: %s (MIME: %s)", filename, mime_type)
        elif mime_type in effective_mime_types:
            # Other allowed files are sent for conversion
            convert_to_pdf.delay(file_path, owner_id=owner_id)
            logger.info("Enqueued file for conversion to PDF: %s", filename)

        has_attachment = True
    return has_attachment


def email_already_has_label(mail, msg_id, label="Ingested"):
    """
    Checks if the given message (msg_id) has the specified Gmail label.
    Returns True if the label is found, False otherwise.
    """
    try:
        # Convert msg_id to bytes if it's an integer
        if isinstance(msg_id, int):
            msg_id = str(msg_id).encode()

        label_status, label_data = mail.fetch(msg_id, "(X-GM-LABELS)")
        if label_status == "OK" and label_data and len(label_data) > 0:
            raw_labels = label_data[0][1].decode("utf-8", errors="ignore")
            if label in raw_labels:
                return True
    except Exception as e:
        logger.error("Failed to fetch labels for msg_id=%s: %s", msg_id, e)
    return False


def mark_as_processed_with_star(mail, msg_id):
    """Stars the email in Gmail."""
    try:
        mail.store(msg_id, "+FLAGS", "\\Flagged")
        logger.info("Email %s starred in Gmail.", msg_id)
    except Exception as e:
        logger.error("Failed to star email %s: %s", msg_id, e)


def mark_as_processed_with_label(mail, msg_id, label="Ingested"):
    """Adds a custom label to the email in Gmail."""
    try:
        mail.store(msg_id, "+X-GM-LABELS", label)
        logger.info("Email %s labeled '%s' in Gmail.", msg_id, label)
    except Exception as e:
        logger.error("Failed to label email %s with %s: %s", msg_id, label, e)


def find_all_mail_folder(mail):
    """
    Attempts to select the Gmail All Mail folder using known localized names.
    Falls back to using XLIST if needed.
    Returns the folder name if found, otherwise None.
    """
    COMMON_ALL_MAIL_NAMES = [
        "[Gmail]/Alle Nachrichten",
        "[Gmail]/All Mail",
        "[Gmail]/Todos",
        "[Gmail]/Tutte le mail",
        "[Gmail]/Tous les messages",
    ]
    for candidate in COMMON_ALL_MAIL_NAMES:
        status, _ = mail.select(f'"{candidate}"', readonly=True)
        if status == "OK":
            return candidate

    capabilities = get_capabilities(mail)
    if "XLIST" in capabilities:
        candidate = find_all_mail_xlist(mail)
        if candidate:
            return candidate
    return None


def get_capabilities(mail):
    """Returns a list of capabilities supported by the IMAP server."""
    typ, data = mail.capability()
    if typ == "OK" and data:
        caps = data[0].decode("utf-8", errors="ignore").upper().split()
        return caps
    return []


def find_all_mail_xlist(mail):
    """
    Uses XLIST to discover the mailbox flagged as All Mail.
    Returns the folder name if found, otherwise None.
    """
    tag = mail._new_tag().decode("ascii")
    command_str = f'{tag} XLIST "" "*"'
    mail.send((command_str + "\r\n").encode("utf-8"))

    all_mail_folder = None
    while True:
        line = mail.readline()
        if not line:
            break
        line_str = line.decode("utf-8", errors="ignore").strip()
        if line_str.upper().startswith("* XLIST ") and "\\ALLMAIL" in line_str.upper():
            match = re.search(r'"([^"]+)"$', line_str)
            if match:
                candidate = match.group(1)
                logger.info("Found All Mail folder via XLIST: %s", candidate)
                all_mail_folder = candidate
        if line_str.startswith(tag):
            break
    return all_mail_folder