"""
Mock Auth API — Test identity linking (device_id ↔ user_id) with in-memory storage.
Ported from chatbot-canifa-dev-experimental for testing merge history UI.

Endpoints:
- POST /api/mock/login        — Link device_id → user_id, return merged history
- POST /api/mock/logout       — Back to device_id mode
- POST /api/mock/history      — Load history (merged if logged in, single if guest)
- POST /api/mock/chat         — Send a mock message
- POST /api/mock/new-conversation — Archive current and start new
- POST /api/mock/conversations — List conversations
- GET  /api/mock/links/{uid}  — View linked device_ids for a user_id
"""

import json
import logging
from collections import defaultdict
from datetime import datetime

from fastapi import APIRouter
from fastapi.responses import JSONResponse
from pydantic import BaseModel

logger = logging.getLogger(__name__)
router = APIRouter(tags=["Mock Auth"])

# ─── In-Memory Storage ──────────────────────────────────────────
# Simple in-memory store for mock testing (no DB dependency)
_identity_links: dict[str, set[str]] = defaultdict(set)  # user_id → {device_ids}
_chat_store: dict[str, list[dict]] = defaultdict(list)   # identity_key → [messages]


# ─── Request Models ──────────────────────────────────────────────

class MockLoginRequest(BaseModel):
    device_id: str
    user_id: str


class MockLogoutRequest(BaseModel):
    device_id: str


class MockChatRequest(BaseModel):
    identity_key: str
    message: str
    use_real_ai: bool = False
    model_name: str | None = None
    images: list[str] | None = None
    is_authenticated: bool = False
    device_id: str | None = None


class MockHistoryRequest(BaseModel):
    identity_key: str
    is_authenticated: bool = False
    limit: int = 50


class MockNewConversationRequest(BaseModel):
    identity_key: str


class MockConversationsRequest(BaseModel):
    identity_key: str
    is_authenticated: bool = False


# ─── Helper Functions ────────────────────────────────────────────

def _now_iso() -> str:
    return datetime.now().isoformat()


def _get_merged_keys(user_id: str) -> list[str]:
    """Get all identity keys linked to a user_id."""
    keys = [user_id]
    keys.extend(_identity_links.get(user_id, set()))
    return list(dict.fromkeys(keys))  # dedupe preserving order


def _get_history(identity_key: str, limit: int = 50) -> list[dict]:
    """Get chat history for a single identity key."""
    return _chat_store.get(identity_key, [])[-limit:]


def _get_merged_history(user_id: str, limit: int = 50) -> list[dict]:
    """Get merged history across all linked identities."""
    all_msgs = []
    for key in _get_merged_keys(user_id):
        all_msgs.extend(_chat_store.get(key, []))
    all_msgs.sort(key=lambda m: m.get("timestamp", ""))
    return all_msgs[-limit:]


# ─── Endpoints ───────────────────────────────────────────────────

@router.post("/api/mock/login", summary="Mock Login — Link device_id to user_id")
async def mock_login(req: MockLoginRequest):
    try:
        _identity_links[req.user_id].add(req.device_id)
        linked = list(_identity_links[req.user_id])
        history = _get_merged_history(req.user_id)

        logger.info(f"🔗 Mock Login: {req.device_id} → {req.user_id} | devices={linked}")
        return {
            "status": "success",
            "identity_key": req.user_id,
            "is_authenticated": True,
            "linked_device_ids": linked,
            "history": history,
        }
    except Exception as e:
        logger.error(f"Mock login error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.post("/api/mock/logout", summary="Mock Logout — Back to device_id")
async def mock_logout(req: MockLogoutRequest):
    try:
        history = _get_history(req.device_id)
        logger.info(f"🔓 Mock Logout: back to {req.device_id}")
        return {
            "status": "success",
            "identity_key": req.device_id,
            "is_authenticated": False,
            "history": history,
        }
    except Exception as e:
        logger.error(f"Mock logout error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.post("/api/mock/history", summary="Load history — merged or single key")
async def mock_get_history(req: MockHistoryRequest):
    try:
        if req.is_authenticated:
            history = _get_merged_history(req.identity_key, req.limit)
        else:
            history = _get_history(req.identity_key, req.limit)

        return {"status": "success", "identity_key": req.identity_key, "history": history}
    except Exception as e:
        logger.error(f"Mock history error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.post("/api/mock/new-conversation", summary="Archive current conversation and start a new one")
async def mock_new_conversation(req: MockNewConversationRequest):
    try:
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        archived_key = f"{req.identity_key}_archived_{ts}"

        # Move current messages to archived key
        if req.identity_key in _chat_store:
            _chat_store[archived_key] = _chat_store.pop(req.identity_key)

        return {
            "status": "success",
            "identity_key": req.identity_key,
            "archived_key": archived_key,
            "started_at": _now_iso(),
        }
    except Exception as e:
        logger.error(f"Mock new conversation error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.post("/api/mock/conversations", summary="List conversations for an identity")
async def mock_conversations(req: MockConversationsRequest):
    try:
        base_keys = [req.identity_key]
        if req.is_authenticated:
            base_keys.extend(_identity_links.get(req.identity_key, set()))
        base_keys = list(dict.fromkeys(base_keys))

        conversations = []
        for key in base_keys:
            # Find current + archived keys
            for store_key, msgs in _chat_store.items():
                if store_key == key or store_key.startswith(f"{key}_archived_"):
                    if msgs:
                        conversations.append({
                            "identity_key": store_key,
                            "message_count": len(msgs),
                            "last_message_at": msgs[-1].get("timestamp", ""),
                            "base_key": key,
                        })

        conversations.sort(key=lambda c: c.get("last_message_at", ""), reverse=True)

        return {
            "status": "success",
            "is_authenticated": req.is_authenticated,
            "identity_key": req.identity_key,
            "conversations": conversations,
        }
    except Exception as e:
        logger.error(f"Mock conversations error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.post("/api/mock/chat", summary="Send a mock chat message")
async def mock_chat(req: MockChatRequest):
    """Save a user message + mock AI reply to in-memory store."""
    try:
        ts = _now_iso()

        # Save human message
        _chat_store[req.identity_key].append({
            "role": "human",
            "message": req.message,
            "timestamp": ts,
            "identity_key": req.identity_key,
        })

        # Generate mock AI response
        if req.use_real_ai:
            # Try to call real AI controller
            try:
                from fastapi import BackgroundTasks
                from agent.controller import chat_controller
                from config import DEFAULT_MODEL

                bg = BackgroundTasks()
                device_id = req.device_id or req.identity_key
                result = await chat_controller(
                    query=req.message,
                    background_tasks=bg,
                    model_name=req.model_name or DEFAULT_MODEL,
                    images=req.images,
                    identity_key=req.identity_key,
                    return_user_insight=False,
                    is_authenticated=req.is_authenticated,
                    device_id=device_id,
                    skip_cache=True,
                )
                ai_response = result.get("ai_response", "")
            except Exception as ai_err:
                logger.warning(f"Real AI failed, falling back to mock: {ai_err}")
                ai_response = f'[Mock AI] Cảm ơn bạn! Tôi đã nhận được tin nhắn: "{req.message}"'
        else:
            ai_response = f'[Mock AI] Cảm ơn bạn! Tôi đã nhận được tin nhắn: "{req.message}"'

        # Save AI response
        _chat_store[req.identity_key].append({
            "role": "ai",
            "message": ai_response,
            "timestamp": _now_iso(),
            "identity_key": req.identity_key,
        })

        return {
            "status": "success",
            "identity_key": req.identity_key,
            "ai_response": ai_response,
            "mode": "real_ai" if req.use_real_ai else "mock_ai",
        }
    except Exception as e:
        logger.error(f"Mock chat error: {e}", exc_info=True)
        return JSONResponse(status_code=500, content={"status": "error", "message": str(e)})


@router.get("/api/mock/links/{user_id}", summary="View linked devices for a user_id")
async def get_links(user_id: str):
    return {
        "user_id": user_id,
        "linked_device_ids": list(_identity_links.get(user_id, set())),
    }
