from __future__ import annotations

from yuxi.channel.message.models import UnifiedMessage
from yuxi.channel.protocols import SessionResolution


class EmailSmtpSessionAdapter:
    def __init__(self):
        self._thread_sessions: dict[str, dict] = {}

    def resolve_session(self, msg: UnifiedMessage, account_id: str = "") -> SessionResolution:
        thread_id = msg.message_thread_id or msg.msg_id
        reply_to_id = msg.reply_to_id
        sender_id = msg.sender.id if msg.sender else "unknown"

        if thread_id not in self._thread_sessions:
            session_id = f"email:{account_id}:{thread_id}"
            self._thread_sessions[thread_id] = {
                "session_id": session_id,
                "root_message_id": thread_id,
                "branch_sessions": {},
            }

        thread = self._thread_sessions[thread_id]

        if reply_to_id and reply_to_id in thread.get("branch_sessions", {}):
            session_id = thread["branch_sessions"][reply_to_id]
        elif reply_to_id and reply_to_id != thread.get("root_message_id"):
            session_id = f"email:{account_id}:{thread_id}:branch:{reply_to_id}"
            thread.setdefault("branch_sessions", {})[reply_to_id] = session_id
        else:
            session_id = thread["session_id"]

        return SessionResolution(
            kind="direct",
            conversation_id=sender_id,
            thread_id=session_id,
            label=msg.metadata.get("subject", "") if msg.metadata else "",
        )

    def extract_thread_id(self, msg: UnifiedMessage) -> str | None:
        return msg.message_thread_id or msg.msg_id