from __future__ import annotations

from typing import TYPE_CHECKING

from yuxi.channels.models import DeliveryResult
from yuxi.utils.logging_config import logger

from .formatter import markdown_to_matrix_html

if TYPE_CHECKING:
    from nio import AsyncClient


class DraftStreamSession:
    def __init__(self, room_id: str):
        self.room_id = room_id
        self.draft_event_id: str = ""
        self.current_text: str = ""
        self.is_active: bool = False

    def start(self, event_id: str) -> None:
        self.draft_event_id = event_id
        self.current_text = ""
        self.is_active = True

    def update(self, text: str) -> None:
        self.current_text = text

    def finish(self) -> None:
        self.is_active = False


# In-memory storage for draft sessions, keyed by room_id.
# TODO: Consider persistent storage for production use.
_DRAFT_SESSIONS: dict[str, DraftStreamSession] = {}


def get_or_create_draft(room_id: str) -> DraftStreamSession:
    if room_id not in _DRAFT_SESSIONS:
        _DRAFT_SESSIONS[room_id] = DraftStreamSession(room_id)
    return _DRAFT_SESSIONS[room_id]


def remove_draft(room_id: str) -> None:
    _DRAFT_SESSIONS.pop(room_id, None)


async def send_draft_update(
    client: AsyncClient,
    room_id: str,
    text: str,
    draft_event_id: str | None = None,
    finished: bool = False,
) -> DeliveryResult:
    session = get_or_create_draft(room_id)
    if draft_event_id and not session.draft_event_id:
        session.start(draft_event_id)

    if not session.is_active:
        resp = await client.room_send(
            room_id=room_id,
            message_type="m.room.message",
            content={
                "msgtype": "m.text",
                "body": text,
                "format": "org.matrix.custom.html",
                "formatted_body": markdown_to_matrix_html(text),
                "org.matrix.msc4357.live": {
                    "draft": True,
                    "finished": finished,
                },
            },
        )
        session.start(resp.event_id)
        return DeliveryResult(success=True, message_id=resp.event_id)

    content = {
        "msgtype": "m.text",
        "body": text,
        "format": "org.matrix.custom.html",
        "formatted_body": markdown_to_matrix_html(text),
        "m.new_content": {
            "msgtype": "m.text",
            "body": text,
            "format": "org.matrix.custom.html",
            "formatted_body": markdown_to_matrix_html(text),
        },
        "m.relates_to": {
            "rel_type": "m.replace",
            "event_id": session.draft_event_id,
        },
    }

    if finished:
        session.finish()
        remove_draft(room_id)

    try:
        resp = await client.room_send(
            room_id=room_id,
            message_type="m.room.message",
            content=content,
        )
        session.update(text)
        return DeliveryResult(success=True, message_id=resp.event_id)
    except Exception as e:
        return DeliveryResult(success=False, error=str(e))


async def send_live_marker(
    client: AsyncClient,
    room_id: str,
    event_id: str,
    content: dict,
    finished: bool = False,
) -> DeliveryResult:
    content.setdefault("org.matrix.msc4357.live", {})
    content["org.matrix.msc4357.live"]["finished"] = finished

    try:
        resp = await client.room_send(
            room_id=room_id,
            message_type="m.room.message",
            content=content,
        )
        return DeliveryResult(success=True, message_id=resp.event_id)
    except Exception as e:
        logger.debug(f"Matrix MSC4357 live marker failed: {e}")
        return DeliveryResult(success=False, error=str(e))