"""
SR 채팅방 API + WebSocket (모바일 기능 #98).

  WS   /ws/sr-chat/{sr_id}?token={jwt}  — SR별 실시간 채팅
  POST /api/sr-chat/{sr_id}/messages    — 메시지 전송 (REST)
  GET  /api/sr-chat/{sr_id}/messages    — 메시지 이력
  POST /api/sr-chat/{sr_id}/read        — 읽음 처리

메시지 타입: text | image | sr_update
"""
from __future__ import annotations

import json
import logging
from datetime import datetime
from typing import Dict, List, Optional, Set

from fastapi import (
    APIRouter, Depends, HTTPException, Query, WebSocket, WebSocketDisconnect,
)
from pydantic import BaseModel, ConfigDict
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from core.auth import get_current_user
from database import get_db, SessionLocal
from models import SRChatMessage, SRRequest, User

logger = logging.getLogger(__name__)
router = APIRouter(tags=["SR Chat"])

_VALID_MSG_TYPE = {"text", "image", "sr_update"}


# ── WebSocket 연결 관리 (SR방별 그룹) ─────────────────────────────────────────
class _ChatRooms:
    def __init__(self) -> None:
        # { sr_id: set(WebSocket) }
        self._rooms: Dict[str, Set[WebSocket]] = {}

    def join(self, sr_id: str, ws: WebSocket) -> None:
        self._rooms.setdefault(sr_id, set()).add(ws)

    def leave(self, sr_id: str, ws: WebSocket) -> None:
        room = self._rooms.get(sr_id)
        if room:
            room.discard(ws)
            if not room:
                self._rooms.pop(sr_id, None)

    async def broadcast(self, sr_id: str, payload: dict) -> None:
        room = self._rooms.get(sr_id)
        if not room:
            return
        msg = json.dumps(payload, ensure_ascii=False)
        dead = []
        for ws in list(room):
            try:
                await ws.send_text(msg)
            except Exception:
                dead.append(ws)
        for ws in dead:
            room.discard(ws)


rooms = _ChatRooms()


# ── 스키마 ────────────────────────────────────────────────────────────────────
class ChatMessageCreate(BaseModel):
    content: str
    msg_type: str = "text"


class ChatMessageOut(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: int
    task_id: str
    sender_id: str
    content: str
    msg_type: str
    created_at: Optional[datetime]


# ── 헬퍼 ──────────────────────────────────────────────────────────────────────
async def _ensure_sr(sr_id: str, db: AsyncSession) -> SRRequest:
    sr = (await db.execute(
        select(SRRequest).where(SRRequest.sr_id == sr_id)
    )).scalars().first()
    if not sr:
        raise HTTPException(404, "SR을 찾을 수 없습니다.")
    return sr


async def _save_message(db: AsyncSession, sr_id: str, sender: str,
                        content: str, msg_type: str) -> SRChatMessage:
    if msg_type not in _VALID_MSG_TYPE:
        raise HTTPException(422, f"msg_type은 {_VALID_MSG_TYPE} 중 하나여야 합니다.")
    if not content or not content.strip():
        raise HTTPException(422, "메시지 내용이 비어 있습니다.")
    m = SRChatMessage(
        task_id=sr_id,
        sender_id=sender,
        content=content,
        msg_type=msg_type,
        read_by=json.dumps([sender], ensure_ascii=False),
    )
    db.add(m)
    await db.commit()
    await db.refresh(m)
    return m


# ── REST: 메시지 전송 ─────────────────────────────────────────────────────────
@router.post("/api/sr-chat/{sr_id}/messages", response_model=ChatMessageOut, status_code=201)
async def send_message(
    sr_id: str,
    payload: ChatMessageCreate,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """SR 채팅 메시지 전송 (REST). 연결된 WebSocket 구독자에게도 브로드캐스트."""
    await _ensure_sr(sr_id, db)
    m = await _save_message(db, sr_id, current_user.username,
                            payload.content, payload.msg_type)
    await rooms.broadcast(sr_id, {
        "type": "message",
        "id": m.id,
        "task_id": sr_id,
        "sender_id": m.sender_id,
        "content": m.content,
        "msg_type": m.msg_type,
        "created_at": m.created_at.isoformat() if m.created_at else None,
    })
    return m


# ── REST: 메시지 이력 ─────────────────────────────────────────────────────────
@router.get("/api/sr-chat/{sr_id}/messages", response_model=List[ChatMessageOut])
async def list_messages(
    sr_id: str,
    skip: int = 0,
    limit: int = 100,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """SR 채팅 메시지 이력 (오래된 순)."""
    await _ensure_sr(sr_id, db)
    rows = (await db.execute(
        select(SRChatMessage)
        .where(SRChatMessage.task_id == sr_id)
        .order_by(SRChatMessage.created_at.asc())
        .offset(skip).limit(min(limit, 500))
    )).scalars().all()
    return rows


# ── REST: 읽음 처리 ───────────────────────────────────────────────────────────
@router.post("/api/sr-chat/{sr_id}/read")
async def mark_read(
    sr_id: str,
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    """현재 사용자가 SR 채팅의 모든 메시지를 읽음 처리."""
    await _ensure_sr(sr_id, db)
    rows = (await db.execute(
        select(SRChatMessage).where(SRChatMessage.task_id == sr_id)
    )).scalars().all()
    updated = 0
    for m in rows:
        try:
            readers = json.loads(m.read_by) if m.read_by else []
        except Exception:
            readers = []
        if current_user.username not in readers:
            readers.append(current_user.username)
            m.read_by = json.dumps(readers, ensure_ascii=False)
            updated += 1
    await db.commit()
    return {"sr_id": sr_id, "marked_read": updated, "reader": current_user.username}


# ── WebSocket: 실시간 채팅 ────────────────────────────────────────────────────
async def _authenticate_ws(token: str, db: AsyncSession) -> Optional[User]:
    if not token:
        return None
    try:
        from core.auth import SECRET_KEY, ALGORITHM
        from jose import jwt
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        if payload.get("mfa_pending"):
            return None
        username = payload.get("sub")
        if not username:
            return None
        user = (await db.execute(
            select(User).where(User.username == username)
        )).scalars().first()
        return user if (user and user.is_active) else None
    except Exception:
        return None


@router.websocket("/ws/sr-chat/{sr_id}")
async def sr_chat_ws(
    websocket: WebSocket,
    sr_id: str,
    token: str = Query(..., description="JWT access_token"),
    db: AsyncSession = Depends(get_db),
):
    """SR별 실시간 채팅 WebSocket."""
    user = await _authenticate_ws(token, db)
    if not user:
        await websocket.close(code=4001, reason="인증 실패: 유효한 토큰이 필요합니다.")
        return

    # SR 존재 확인
    sr = (await db.execute(
        select(SRRequest).where(SRRequest.sr_id == sr_id)
    )).scalars().first()
    if not sr:
        await websocket.close(code=4004, reason="SR을 찾을 수 없습니다.")
        return

    await websocket.accept()
    rooms.join(sr_id, websocket)
    await websocket.send_text(json.dumps({
        "type": "connected",
        "sr_id": sr_id,
        "username": user.username,
        "server_time": datetime.now().isoformat(),
    }, ensure_ascii=False))

    try:
        while True:
            raw = await websocket.receive_text()
            try:
                data = json.loads(raw)
            except Exception:
                await websocket.send_text(json.dumps(
                    {"type": "error", "message": "JSON 형식이 아닙니다."},
                    ensure_ascii=False))
                continue

            if data.get("type") == "ping":
                await websocket.send_text(json.dumps(
                    {"type": "pong", "server_time": datetime.now().isoformat()},
                    ensure_ascii=False))
                continue

            content = (data.get("content") or "").strip()
            msg_type = data.get("msg_type", "text")
            if not content or msg_type not in _VALID_MSG_TYPE:
                await websocket.send_text(json.dumps(
                    {"type": "error", "message": "content 또는 msg_type이 올바르지 않습니다."},
                    ensure_ascii=False))
                continue

            # DB 저장 (독립 세션) + 구독자에게 브로드캐스트
            async with SessionLocal() as _db:
                m = await _save_message(_db, sr_id, user.username, content, msg_type)
                payload = {
                    "type": "message",
                    "id": m.id,
                    "task_id": sr_id,
                    "sender_id": m.sender_id,
                    "content": m.content,
                    "msg_type": m.msg_type,
                    "created_at": m.created_at.isoformat() if m.created_at else None,
                }
            await rooms.broadcast(sr_id, payload)

    except WebSocketDisconnect:
        pass
    except Exception as exc:
        logger.debug("SR 채팅 WS 오류: sr=%s err=%s", sr_id, exc)
    finally:
        rooms.leave(sr_id, websocket)
