"""IM channel - connects to SuperAgent IM via WebSocket long connection."""

from __future__ import annotations

import asyncio
import json
import mimetypes
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

import httpx
import websockets
from loguru import logger

from asedclaw.bus.events import InboundAttachment, OutboundMessage
from asedclaw.bus.queue import MessageBus
from asedclaw.channels.base import BaseChannel
from asedclaw.config.paths import get_media_dir
from asedclaw.config.schema import ImConfig
from asedclaw.media.fetch import download_to_media_store
from asedclaw.media.store import MediaStore


@dataclass
class _StreamState:
    """Tracks the state of an in-progress streaming response for one chat."""
    stream_id: str
    sequence: int = 0
    accumulated: str = ""
    typing_sent: bool = False
    task_id: str | None = None  # agent task identifier
    item_type: str | None = None  # answer_text | reasoning_content


@dataclass
class _UploadedFile:
    """Normalized IM file metadata returned by the upload API."""

    file_id: str
    file_name: str | None = None
    file_size: int | None = None
    mime_type: str | None = None
    file_url: str | None = None


def _attachment_kind(mime_type: str | None) -> str:
    if not mime_type:
        return "file"
    major = mime_type.split("/", 1)[0]
    if major in {"image", "audio", "video", "text"}:
        return major
    if mime_type == "application/pdf":
        return "document"
    return "file"


class IMChannel(BaseChannel):
    """IM channel connecting to SuperAgent IM service via WebSocket long connection."""

    name = "ased"

    def __init__(self, config: ImConfig, bus: MessageBus):
        super().__init__(config, bus)
        self.config: ImConfig = config
        self._token: str | None = None
        self._token_expires_at: float = 0
        self._ws: websockets.WebSocketClientProtocol | None = None
        self._refresh_margin = 300  # Refresh token 5 min before expiry
        self._conv_to_user: dict[str, str] = {}  # conversation_id -> user_id
        self._keepalive_task: asyncio.Task[None] | None = None
        # stream_id per chat_id, cleared when a non-progress message arrives
        self._streams: dict[str, _StreamState] = {}
        self._media_store: MediaStore | None = None

    # ──────────────────────────────────────────────
    # Auth
    # ──────────────────────────────────────────────

    def is_allowed(self, sender_id: str) -> bool:
        """IM: empty or ['*'] = allow all; specific ids = allowlist."""
        allow_list = getattr(self.config, "allow_from", []) or []
        if not allow_list or "*" in allow_list:
            return True
        sender_str = str(sender_id)
        return sender_str in allow_list or any(
            p in allow_list for p in sender_str.split("|") if p
        )

    async def _get_token(self) -> str | None:
        """Get valid access token, refreshing if needed."""
        now = time.time()
        if self._token and self._token_expires_at > now + self._refresh_margin:
            return self._token
        base = self.config.im_base_url.rstrip("/")
        url = f"{base}/api/im/v1/app/token"
        try:
            async with httpx.AsyncClient(timeout=15.0) as client:
                resp = await client.post(
                    url,
                    json={
                        "app_id": self.config.app_id,
                        "app_secret": self.config.app_secret,
                    },
                )
                if resp.status_code != 200:
                    logger.error("IM token fetch failed: {}", resp.text)
                    return None
                data = resp.json()
                if not data.get("status") or not data.get("data"):
                    return None
                tok = data["data"].get("access_token")
                exp = data["data"].get("expires_in", 7200)
                if tok:
                    self._token = tok
                    self._token_expires_at = now + exp
                    return self._token
        except Exception as e:
            logger.error("IM token fetch error: {}", e)
        return None

    # ──────────────────────────────────────────────
    # WebSocket receive loop
    # ──────────────────────────────────────────────

    async def _run_websocket(self) -> None:
        """WebSocket receive loop with automatic reconnect."""
        base = self.config.im_base_url.rstrip("/")
        ws_url = base.replace("https://", "wss://").replace("http://", "ws://")
        ws_url = f"{ws_url}/api/im/ws/bot"
        while self._running:
            token = await self._get_token()
            if not token:
                await asyncio.sleep(10)
                continue
            try:
                async with websockets.connect(
                    f"{ws_url}?token={token}",
                    ping_interval=25,
                    ping_timeout=30,
                    close_timeout=5,
                ) as ws:
                    self._ws = ws
                    logger.info("IM WebSocket connected")
                    self._start_keepalive(ws)
                    try:
                        while self._running:
                            try:
                                raw = await asyncio.wait_for(ws.recv(), timeout=35)
                                msg = json.loads(raw)
                                logger.debug("IM WebSocket message: {}", msg)
                                await self._handle_ws_message(msg)
                            except asyncio.TimeoutError:
                                continue
                            except websockets.ConnectionClosed:
                                logger.info("IM WebSocket connection closed")
                                break
                    finally:
                        await self._stop_keepalive()
            except Exception as e:
                logger.warning("IM WebSocket error: {}", e)
            finally:
                self._ws = None
            if self._running:
                await asyncio.sleep(5)

    def _start_keepalive(self, ws: Any) -> None:
        """Start application-level keepalive ping."""
        interval = getattr(self.config, "ws_keepalive_interval_s", 0) or 0
        if interval <= 0:
            return

        async def loop() -> None:
            while self._running and self._ws is ws:
                await asyncio.sleep(interval)
                if not self._running or self._ws is not ws:
                    return
                try:
                    await ws.send(json.dumps({"type": "ping"}))
                    logger.debug("IM WebSocket keepalive ping sent")
                except websockets.ConnectionClosed:
                    break
                except Exception as e:
                    logger.warning("IM WebSocket keepalive error: {}", e)

        self._keepalive_task = asyncio.create_task(loop())

    async def _stop_keepalive(self) -> None:
        if self._keepalive_task:
            t = self._keepalive_task
            self._keepalive_task = None
            t.cancel()
            try:
                await t
            except asyncio.CancelledError:
                pass

    # ──────────────────────────────────────────────
    # Inbound message handling
    # ──────────────────────────────────────────────

    async def _handle_ws_message(self, msg: dict) -> None:
        """Handle inbound messages from WebSocket."""
        t = msg.get("type")
        if t == "message.new":
            await self._emit_inbound(msg.get("data", {}))
        elif t == "messages.offline":
            for m in (msg.get("data") or {}).get("messages") or []:
                await self._emit_inbound(m)
        # pong / message.ack / unknown types are silently ignored

    async def _emit_inbound(self, data: dict) -> None:
        """Convert IM message to InboundMessage and publish to the bus."""
        sender_id = str(data.get("user_id", ""))
        chat_id = str(data.get("conversation_id", ""))
        self._conv_to_user[chat_id] = sender_id
        msg_type = data.get("msg_type", "text")
        content_obj = data.get("content") or {}
        if msg_type == "text":
            content = content_obj.get("text", "")
        else:
            raw_text = content_obj.get("text")
            content = raw_text.strip() if isinstance(raw_text, str) else ""
        attachments = await self._build_inbound_attachments(msg_type, content_obj)
        media = [a.path for a in attachments if a.path] or [a.url for a in attachments if a.url]
        session_key = f"ased:{self.config.app_id}:{chat_id}"

        # Send typing.start immediately when we receive the user's message,
        # so the user sees a feedback before the LLM even starts generating.
        await self._ws_send({
            "type": "message.typing",
            "data": {"user_id": int(sender_id) if str(sender_id).isdigit() else sender_id, "status": "start"},
        })

        await self._handle_message(
            sender_id=sender_id,
            chat_id=chat_id,
            content=content,
            media=media if media else None,
            attachments=attachments or None,
            metadata={
                "im_msg_type": msg_type,
                "im_content": content_obj,
                "attachments": [self._attachment_metadata(a) for a in attachments],
            },
            session_key=session_key,
        )

    async def _build_inbound_attachments(
        self,
        msg_type: str,
        content_obj: dict[str, Any],
    ) -> list[InboundAttachment]:
        if msg_type != "file":
            return []

        file_id = content_obj.get("file_id")
        if isinstance(file_id, str) and file_id.strip():
            file_id = file_id.strip()
        else:
            file_id = None

        file_url = self._resolve_inbound_file_url(content_obj)
        if not file_id and not file_url:
            return []
        file_name = content_obj.get("file_name")
        if not isinstance(file_name, str) or not file_name.strip():
            file_name = "file"
        file_name = file_name.strip()

        size = content_obj.get("file_size")
        declared_size = size if isinstance(size, int) and size >= 0 else None
        mime_type = content_obj.get("mime_type")
        if not isinstance(mime_type, str) or not mime_type.strip():
            mime_type = mimetypes.guess_type(file_name)[0]

        attachment = InboundAttachment(
            file_id=file_id,
            url=file_url,
            name=file_name,
            mime_type=mime_type,
            size=declared_size,
            kind=_attachment_kind(mime_type),
            source=self.name,
        )
        try:
            token = await self._get_token()
            if not token:
                raise RuntimeError("missing app token for inbound file download")
            download_url = (
                self._build_file_download_url(file_id)
                if file_id
                else file_url
            )
            if not download_url:
                raise RuntimeError("missing file download url")
            media_store = self._get_media_store()
            downloaded = await download_to_media_store(
                download_url,
                media_store,
                filename=file_name,
                timeout_s=self.config.inbound_file_timeout_s,
                max_bytes=self.config.inbound_file_max_bytes,
                headers={"Authorization": f"Bearer {token}"},
            )
            attachment.path = downloaded.path
            attachment.size = downloaded.size
            attachment.mime_type = downloaded.mime_type or attachment.mime_type
            attachment.kind = _attachment_kind(attachment.mime_type)
        except Exception as exc:
            logger.warning(
                "ASED inbound attachment download failed for file_id={} url={}: {}",
                file_id,
                file_url,
                exc,
            )
        return [attachment]

    def _get_media_store(self) -> MediaStore:
        if self._media_store is None:
            ttl = max(int(self.config.inbound_file_retention_hours), 1) * 3600
            self._media_store = MediaStore(get_media_dir("ased"), ttl_seconds=ttl)
        return self._media_store

    def _build_file_download_url(self, file_id: str | None) -> str:
        if not file_id:
            raise ValueError("missing file_id for ASED file download")
        base = self.config.im_base_url.rstrip("/")
        return f"{base}/api/im/v1/files/{file_id}"

    def _resolve_inbound_file_url(self, content_obj: dict[str, Any]) -> str | None:
        file_url = content_obj.get("file_url")
        if isinstance(file_url, str) and file_url.strip():
            return file_url.strip()
        file_id = content_obj.get("file_id")
        if isinstance(file_id, str) and file_id.strip():
            return self._build_file_download_url(file_id.strip())
        return None

    @staticmethod
    def _attachment_metadata(attachment: InboundAttachment) -> dict[str, Any]:
        return {
            "file_id": attachment.file_id,
            "path": attachment.path,
            "url": attachment.url,
            "name": attachment.name,
            "mime_type": attachment.mime_type,
            "size": attachment.size,
            "kind": attachment.kind,
            "source": attachment.source,
        }

    # ──────────────────────────────────────────────
    # Lifecycle
    # ──────────────────────────────────────────────

    async def start(self) -> None:
        """Start IM channel WebSocket connection."""
        if not self.config.app_id or not self.config.app_secret:
            logger.error("IM channel: app_id or app_secret not configured")
            return
        self._running = True
        asyncio.create_task(self._run_websocket())

    async def stop(self) -> None:
        """Stop IM channel."""
        self._running = False
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None
        self._streams.clear()
        logger.info("IM channel stopped")

    # ──────────────────────────────────────────────
    # Outbound: send via Bot WebSocket
    # ──────────────────────────────────────────────

    async def _ws_send(self, payload: dict) -> bool:
        """Send a JSON frame on the Bot WebSocket. Returns False if not connected.

        websockets >= 14 uses ClientConnection which has no .closed attribute;
        rely on exception handling instead.
        """
        ws = self._ws
        if ws is None:
            return False
        try:
            await ws.send(json.dumps(payload))
            return True
        except websockets.ConnectionClosed:
            logger.warning("IM WS send failed: connection closed")
            return False
        except Exception as e:
            logger.warning("IM WS send error: {}", e)
            return False

    async def send(self, msg: OutboundMessage) -> None:
        """Route outbound message over the Bot WebSocket.

        Routing logic:
        - _progress=True, _tool_hint=False → stream text chunk (message.stream)
        - _progress=True, _tool_hint=True  → tool hint card (message.send, msg_type=tool),
                                             does NOT end the current stream
        - normal message                   → finalize any active stream, then message.send
        """
        chat_id = msg.chat_id
        user_id = self._user_id_from_chat(chat_id)
        if not user_id:
            logger.warning("IM send: cannot derive user_id from chat_id {}", chat_id)
            return

        is_progress: bool = bool(msg.metadata.get("_progress"))
        is_tool_hint: bool = bool(msg.metadata.get("_tool_hint"))
        task_id: str | None = msg.metadata.get("_task_id") or msg.task_id
        item_type: str = str(msg.metadata.get("_item_type") or "answer_text")

        if is_progress and not is_tool_hint:
            await self._send_stream_chunk(
                chat_id,
                user_id,
                msg.content or "",
                task_id=task_id,
                item_type=item_type,
            )
        elif is_progress and is_tool_hint:
            # Keep history order closer to realtime: flush current reasoning segment
            # before inserting a tool event, so tool/reasoning can interleave after reload.
            current = self._streams.get(chat_id)
            if current and current.item_type == "reasoning_content":
                await self._finalize_stream(chat_id, user_id)
            # Tool call hint during streaming: send as a tool card without ending the stream
            await self._send_tool_hint(chat_id, user_id, msg, task_id=task_id)
        else:
            # If an active stream existed, _finalize_stream already persists the full text
            # via message.stream.end on the server side → skip _send_normal to avoid duplicate.
            stream_was_active = await self._finalize_stream(chat_id, user_id)
            if not stream_was_active:
                await self._send_normal(
                    chat_id,
                    user_id,
                    msg,
                    task_id=task_id,
                    item_type=item_type,
                )

    async def _send_stream_chunk(
        self,
        chat_id: str,
        user_id: str,
        delta: str,
        *,
        task_id: str | None = None,
        item_type: str | None = None,
    ) -> None:
        """Send a streaming text chunk. Opens the stream with typing.start if needed."""
        state = self._streams.get(chat_id)
        if state is not None and item_type and state.item_type and state.item_type != item_type:
            await self._finalize_stream(chat_id, user_id)
            state = None
        if state is None:
            state = _StreamState(stream_id=uuid.uuid4().hex, task_id=task_id, item_type=item_type)
            self._streams[chat_id] = state

        if not state.typing_sent:
            # typing.start was already sent in _emit_inbound when the user message arrived.
            # Mark as sent to avoid sending again; no need to send here.
            state.typing_sent = True

        state.sequence += 1
        state.accumulated += delta
        chunk_data: dict = {
            "user_id": int(user_id),
            "stream_id": state.stream_id,
            "delta": delta,
            "sequence": state.sequence,
        }
        if state.task_id:
            chunk_data["task_id"] = state.task_id
        if state.item_type:
            chunk_data["item_type"] = state.item_type
        await self._ws_send({"type": "message.stream", "data": chunk_data})

    async def _finalize_stream(self, chat_id: str, user_id: str) -> bool:
        """End an active stream: send stream.end + typing.stop, then clear state.

        Returns True if a stream was active (and thus the message has already been
        persisted by the server via message.stream.end).
        Returns False if no stream was active.
        """
        state = self._streams.pop(chat_id, None)
        if state is None:
            return False
        end_data: dict = {
            "user_id": int(user_id),
            "stream_id": state.stream_id,
            "full_text": state.accumulated,
        }
        if state.task_id:
            end_data["task_id"] = state.task_id
        if state.item_type:
            end_data["item_type"] = state.item_type
        await self._ws_send({"type": "message.stream.end", "data": end_data})
        if state.typing_sent:
            await self._ws_send({
                "type": "message.typing",
                "data": {"user_id": int(user_id), "status": "stop"},
            })
        return True

    async def _send_tool_hint(self, chat_id: str, user_id: str, msg: OutboundMessage, *, task_id: str | None = None) -> None:
        """Send or update a tool call hint card during streaming.

        Lifecycle:
        - tool_call_started  → status="running", no result  (new card)
        - tool_call_completed → status="done", result set   (update card via tool_id)
        """
        meta = msg.metadata or {}
        # tool_name: prefer _tool_name metadata, fall back to the formatted hint string
        tool_name = meta.get("_tool_name") or msg.content or "unknown"
        tool_input = meta.get("_tool_input") or {}
        tool_result = meta.get("_tool_result")          # None → running
        status = meta.get("_tool_status") or ("done" if tool_result is not None else "running")
        tool_id = meta.get("_tool_id")                  # links completion back to start card

        # Parse structured tool name + input from the formatted hint string when not explicit.
        # Hint format: 'read_file("/some/path…")' or just 'read_file'
        if not meta.get("_tool_name") and msg.content:
            import re as _re
            m = _re.match(r'^(\w+)\("(.*)"\)$', msg.content)
            if m:
                tool_name = m.group(1)

        content: dict = {
            "tool_name": tool_name,
            "tool_input": tool_input,
            "tool_result": tool_result,
            "status": status,
            "item_type": "tool_call",
        }
        if tool_id:
            content["tool_id"] = tool_id
        if task_id:
            content["task_id"] = task_id

        await self._ws_send({
            "type": "message.send",
            "data": {
                "user_id": int(user_id),
                "msg_type": "tool",
                "content": content,
            },
        })
        logger.debug("IM tool hint: tool={} status={} id={}", tool_name, status, tool_id)

    async def _upload_file(self, local_path: str) -> _UploadedFile | None:
        """Upload a local file to the IM server and return normalized metadata.

        Uses the same app token as the Bot WS connection.
        Returns None on failure.
        """
        from pathlib import Path as _Path
        import httpx as _httpx

        token = await self._get_token()
        if not token:
            logger.warning("IM upload: no token available")
            return None

        file_path = _Path(local_path)
        if not file_path.is_file():
            logger.warning("IM upload: file not found: {}", local_path)
            return None

        base = self.config.im_base_url.rstrip("/")
        url = f"{base}/api/im/v1/files/upload"

        try:
            async with _httpx.AsyncClient(timeout=60.0) as client:
                with file_path.open("rb") as fh:
                    resp = await client.post(
                        url,
                        files={"file": (file_path.name, fh)},
                        headers={"Authorization": f"Bearer {token}"},
                    )
                if resp.status_code != 200:
                    logger.error(
                        "IM upload failed ({}): {}", resp.status_code, resp.text
                    )
                    return None
                data = resp.json()
                if not data.get("status") or not data.get("data"):
                    logger.error("IM upload bad response: {}", data)
                    return None
                uploaded = self._normalize_uploaded_file(data.get("data"), fallback_name=file_path.name)
                if uploaded:
                    logger.debug(
                        "IM upload OK: {} → file_id={}", file_path.name, uploaded.file_id
                    )
                return uploaded
        except Exception as e:
            logger.error("IM upload error for {}: {}", local_path, e)
            return None

    def _normalize_uploaded_file(
        self,
        data: dict[str, Any] | None,
        *,
        fallback_name: str | None = None,
    ) -> _UploadedFile | None:
        """Normalize file upload response into the shape expected by the frontend."""

        if not isinstance(data, dict):
            return None

        file_id = data.get("file_id")
        if not isinstance(file_id, str) or not file_id.strip():
            return None

        file_name = data.get("file_name")
        if not isinstance(file_name, str) or not file_name.strip():
            file_name = fallback_name

        file_size = data.get("file_size")
        if not isinstance(file_size, int):
            file_size = None

        mime_type = data.get("mime_type")
        if not isinstance(mime_type, str) or not mime_type.strip():
            mime_type = None

        file_url = data.get("file_url")
        if not isinstance(file_url, str) or not file_url.strip():
            file_url = None

        return _UploadedFile(
            file_id=file_id,
            file_name=file_name,
            file_size=file_size,
            mime_type=mime_type,
            file_url=file_url,
        )

    def _build_file_content(
        self,
        file_id: str,
        uploaded: _UploadedFile | None = None,
    ) -> dict[str, Any]:
        """Build the outbound file payload expected by the ASED IM frontend."""

        content: dict[str, Any] = {"file_id": file_id}
        if uploaded is None:
            return content
        if uploaded.file_name:
            content["file_name"] = uploaded.file_name
        if uploaded.file_size is not None:
            content["file_size"] = uploaded.file_size
        if uploaded.mime_type:
            content["mime_type"] = uploaded.mime_type
        if uploaded.file_url:
            content["file_url"] = uploaded.file_url
        return content

    async def _send_normal(
        self,
        chat_id: str,
        user_id: str,
        msg: OutboundMessage,
        *,
        task_id: str | None = None,
        item_type: str | None = None,
    ) -> None:
        """Send a regular (non-streaming) message via message.send.

        Media handling:
        - If media[0] is a UUID-like file_id → send directly as msg_type=file.
        - If media[0] is a local file path → upload to IM server first, then send.
        - Uploaded file metadata is forwarded in `content` so the frontend renders
          the same file card shape as page-sent messages.
        - Multiple media items are each uploaded and sent as separate file messages.
        - If content is empty and no media succeeds, nothing is sent.
        """
        media_items: list[str] = list(msg.media or [])
        text_content: str = msg.content or ""

        # ── Send text (if any) ──────────────────────────────────────────
        if text_content:
            text_payload: dict = {"text": text_content}
            if task_id:
                text_payload["_task_id"] = task_id
            if item_type:
                text_payload["item_type"] = item_type
            await self._ws_send({
                "type": "message.send",
                "data": {
                    "user_id": int(user_id),
                    "msg_type": "text",
                    "content": text_payload,
                },
            })

        # ── Send each media item ─────────────────────────────────────────
        for item in media_items:
            file_id: str | None = None
            uploaded: _UploadedFile | None = None

            if len(item) == 36 and "-" in item:
                # Already a UUID file_id (e.g. from a previous upload)
                file_id = item
            else:
                # Local file path — upload to IM server first
                uploaded = await self._upload_file(item)
                file_id = uploaded.file_id if uploaded else None

            if not file_id:
                logger.warning("IM send_normal: skipping media item (upload failed): {}", item)
                continue

            file_content = self._build_file_content(file_id, uploaded)
            if task_id:
                file_content["_task_id"] = task_id
            ok = await self._ws_send({
                "type": "message.send",
                "data": {
                    "user_id": int(user_id),
                    "msg_type": "file",
                    "content": file_content,
                },
            })
            if ok:
                logger.debug("IM file sent to user {}: file_id={}", user_id, file_id)
            else:
                logger.warning("IM WS unavailable, file to user {} dropped", user_id)

        # ── If no text and no media, send a plain text fallback ─────────
        if not text_content and not media_items:
            logger.debug("IM send_normal: nothing to send for user {}", user_id)

    # ──────────────────────────────────────────────
    # Helpers
    # ──────────────────────────────────────────────

    def _user_id_from_chat(self, chat_id: str) -> str | None:
        """Derive user_id from conversation_id (stored when receiving messages)."""
        return self._conv_to_user.get(chat_id)
