"""
app/utils/sse_utils.py — SSE 工具函数（按 Lead 规范路径）
补发历史事件、event_seq 去重、Redis pub/sub 推送。
供 api/v1/stream.py 和 AIE worker 共用。
"""

import asyncio
import json
import logging
from typing import Any, AsyncGenerator

from app.core.config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()

# SSE 事件类型（契约§2全列）
SSE_EVENT_TYPES = frozenset({
    "task_started", "analyze_done", "text_progress", "text_candidate",
    "image_progress", "image_candidate", "flywheel_injected",
    "batch_failed", "task_done", "error", "heartbeat",
})

_HISTORY_KEY_TPL = "sse:task:{task_id}:events"
_CHANNEL_TPL = "sse:task:{task_id}"
_HISTORY_TTL_SECONDS = 3600  # 历史事件保留 1h


def format_sse(event: str, data: dict, seq: int | None = None) -> str:
    """格式化单条 SSE 消息（text/event-stream 格式）。"""
    payload = json.dumps(data, ensure_ascii=False)
    lines = [f"event: {event}", f"data: {payload}"]
    if seq is not None:
        lines.append(f"id: {seq}")
    return "\n".join(lines) + "\n\n"


async def push_event(
    redis_client,
    task_id: int,
    workspace_id: int,
    event: str,
    data: dict[str, Any],
    seq: int,
) -> None:
    """
    推送事件到 Redis：
    1. 追加到历史 list（供断线重连补发）
    2. Publish 到 channel（供在线客户端实时收）
    """
    if event not in SSE_EVENT_TYPES:
        logger.warning("Unknown SSE event type: %s", event)

    record = {"event": event, "data": data, "seq": seq, "workspace_id": workspace_id}
    payload = json.dumps(record, ensure_ascii=False)

    hist_key = _HISTORY_KEY_TPL.format(task_id=task_id)
    channel = _CHANNEL_TPL.format(task_id=task_id)

    await redis_client.rpush(hist_key, payload)
    await redis_client.expire(hist_key, _HISTORY_TTL_SECONDS)
    await redis_client.publish(channel, payload)


async def get_history_events(
    redis_client, task_id: int, after_seq: int
) -> list[dict]:
    """取 task 历史事件中 seq > after_seq 的部分（断线重连补发）。"""
    hist_key = _HISTORY_KEY_TPL.format(task_id=task_id)
    try:
        raw_list = await redis_client.lrange(hist_key, 0, -1)
        result = []
        for raw in raw_list:
            ev = json.loads(raw)
            if ev.get("seq", 0) > after_seq:
                result.append(ev)
        return result
    except Exception as exc:
        logger.warning("get_history_events failed task=%s: %s", task_id, exc)
        return []


async def stream_events(
    redis_client,
    task_id: int,
    workspace_id: int,
    last_seq: int = 0,
    heartbeat_interval: int = 25,
) -> AsyncGenerator[str, None]:
    """
    主 SSE 生成器：
    1. 先补发 last_seq 之后的历史事件
    2. 再订阅 Redis channel 实时推新事件
    3. 每 heartbeat_interval 秒推一次保活 heartbeat
    """
    # 1. 补发历史
    history = await get_history_events(redis_client, task_id, last_seq)
    for ev in history:
        yield format_sse(ev["event"], ev["data"], ev.get("seq"))

    # 2. 实时订阅
    pubsub = redis_client.pubsub()
    channel = _CHANNEL_TPL.format(task_id=task_id)
    await pubsub.subscribe(channel)
    elapsed = 0
    try:
        while True:
            msg = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
            if msg and msg["type"] == "message":
                ev = json.loads(msg["data"])
                if ev.get("workspace_id") != workspace_id:
                    continue  # 防越权订阅他人任务
                yield format_sse(ev["event"], ev["data"], ev.get("seq"))
                if ev["event"] in ("task_done", "error"):
                    break
            else:
                elapsed += 1
                if elapsed >= heartbeat_interval:
                    elapsed = 0
                    yield format_sse("heartbeat", {"ts": asyncio.get_event_loop().time()})
    finally:
        await pubsub.unsubscribe(channel)