import asyncio
import ipaddress
import logging
from datetime import datetime, timezone

from aiosmtpd.smtp import SMTP
from sqlalchemy import func, select

from app.config import get_settings
from app.database import SessionFactory, mongo_db
from app.models import Account
from app.services.mime_parser import parse_mime
from app.services.raw_mime_storage import delivery_digest, get_raw_mime_storage

logger = logging.getLogger(__name__)


def peer_is_allowed(session) -> bool:
    if session is None or not getattr(session, "peer", None):
        return True
    try:
        peer = ipaddress.ip_address(session.peer[0])
        return any(
            peer in ipaddress.ip_network(cidr)
            for cidr in get_settings().inbound_allowed_cidrs
        )
    except ValueError:
        return False


class InboundHandler:
    async def handle_RCPT(self, server, session, envelope, address, rcpt_options):
        if not peer_is_allowed(session):
            logger.warning("Rejected inbound SMTP peer %s", session.peer[0])
            return "550 5.7.1 Delivery from this host is not authorized"
        normalized = address.strip().lower()
        async with SessionFactory() as database_session:
            account = await database_session.scalar(
                select(Account).where(func.lower(Account.email) == normalized)
            )
        if account is None:
            return "550 5.1.1 Recipient address rejected: user unknown"
        envelope.rcpt_tos.append(normalized)
        envelope.rcpt_options.extend(rcpt_options)
        if not hasattr(envelope, "account_ids"):
            envelope.account_ids = {}
        envelope.account_ids[normalized] = str(account.id)
        return "250 2.1.5 Recipient OK"

    async def handle_DATA(self, server, session, envelope):
        settings = get_settings()
        raw = envelope.original_content
        if len(raw) > settings.inbound_max_bytes:
            return "552 5.3.4 Message size exceeds fixed maximum message size"
        recipients = list(dict.fromkeys(address.lower() for address in envelope.rcpt_tos))
        if not recipients:
            return "554 5.5.1 No valid recipients"

        delivery_key = delivery_digest(raw, recipients)
        try:
            tombstones = getattr(mongo_db, "message_tombstones", None)
            if tombstones is not None and await tombstones.find_one(
                {"delivery_key": delivery_key},
                {"_id": 1},
            ):
                return f"250 2.0.0 Message previously accepted as {delivery_key}"
            parsed = parse_mime(raw)
            body_text = parsed.pop("body_text")
            body_html = parsed.pop("body_html")
            stored = await get_raw_mime_storage().put(raw, delivery_key)
            now = datetime.now(timezone.utc)
            message = {
                **parsed,
                "direction": "inbound",
                "envelope_sender": envelope.mail_from.lower(),
                "envelope_recipients": recipients,
                "delivery_key": delivery_key,
                "raw_mime": {
                    "bucket": stored.bucket,
                    "name": stored.name,
                    "generation": stored.generation,
                    "size": stored.size,
                    "crc32c": stored.crc32c,
                },
                "received_at": now,
                "created_at": now,
                "has_text_body": bool(body_text),
                "has_html_body": bool(body_html),
            }
            await mongo_db.messages.update_one(
                {"delivery_key": delivery_key},
                {"$setOnInsert": message},
                upsert=True,
            )
            persisted = await mongo_db.messages.find_one(
                {"delivery_key": delivery_key},
                {"_id": 1},
            )
            if persisted is None:
                raise RuntimeError("Inbound message was not persisted.")
            for recipient in recipients:
                account_id = envelope.account_ids[recipient]
                await mongo_db.mailbox_entries.update_one(
                    {
                        "account_id": account_id,
                        "message_id": persisted["_id"],
                        "mailbox": "inbox",
                    },
                    {
                        "$setOnInsert": {
                            "account_id": account_id,
                            "message_id": persisted["_id"],
                            "mailbox": "inbox",
                            "direction": "inbound",
                            "sender": parsed["sender"],
                            "recipients": parsed["to"],
                            "subject": parsed["subject"],
                            "snippet": parsed["snippet"],
                            "attachments": parsed["attachments"],
                            "status": "received",
                            "is_read": False,
                            "is_starred": False,
                            "event_at": now,
                            "created_at": now,
                            "trashed_at": None,
                            "deletion_state": None,
                        }
                    },
                    upsert=True,
                )
        except Exception:
            logger.exception("Temporary failure accepting inbound message")
            return "451 4.3.0 Temporary server failure; please retry"
        return f"250 2.0.0 Message accepted as {delivery_key}"


class InboundSMTPServer:
    def __init__(self):
        self.server: asyncio.Server | None = None

    async def start(self) -> None:
        settings = get_settings()
        loop = asyncio.get_running_loop()
        self.server = await loop.create_server(
            lambda: SMTP(
                InboundHandler(),
                data_size_limit=settings.inbound_max_bytes,
                enable_SMTPUTF8=True,
                decode_data=False,
            ),
            host=settings.inbound_smtp_host,
            port=settings.inbound_smtp_port,
        )

    async def stop(self) -> None:
        if self.server is not None:
            self.server.close()
            await self.server.wait_closed()
            self.server = None