import asyncio
import hmac
import json
import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, Request, WebSocket, WebSocketDisconnect, Security
from fastapi.security import APIKeyHeader, APIKeyQuery
from pydantic import BaseModel

from agentkit.core.protocol import TaskMessage
from agentkit.core.react import ReActEngine
from agentkit.router.intent import IntentRouter

logger = logging.getLogger(__name__)

router = APIRouter(tags=["portal"])

# ---------------------------------------------------------------------------
# API Key Authentication
# ---------------------------------------------------------------------------

_api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
_api_key_query = APIKeyQuery(name="api_key", auto_error=False)


async def _verify_api_key(
    request: Request,
    api_key_header: str | None = Security(_api_key_header),
    api_key_query: str | None = Security(_api_key_query),
) -> None:
    """Verify API key for REST endpoints. Raises HTTPException if invalid."""
    configured_api_key: str | None = None
    if hasattr(request.app.state, "server_config") and request.app.state.server_config:
        configured_api_key = request.app.state.server_config.api_key
    if configured_api_key is None and hasattr(request.app.state, "api_key"):
        configured_api_key = request.app.state.api_key

    # If no API key is configured, allow all requests (backwards compat)
    if configured_api_key is None:
        return

    provided = api_key_header or api_key_query
    if not hmac.compare_digest((provided or "").encode(), configured_api_key.encode()):
        raise HTTPException(
            status_code=401,
            detail="Invalid or missing API key. Provide via X-API-Key header or api_key query parameter.",
        )


# ---------------------------------------------------------------------------
# In-memory Conversation Store
# ---------------------------------------------------------------------------


@dataclass
class ChatMessage:
    role: str  # "user" or "assistant"
    content: str
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: dict = field(default_factory=dict)


@dataclass
class Conversation:
    id: str
    messages: list[ChatMessage] = field(default_factory=list)
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


class ConversationStore:
    def __init__(self, max_conversations: int = 1000):
        self._conversations: dict[str, Conversation] = {}
        self._max = max_conversations

    def get_or_create(self, conversation_id: str | None = None) -> Conversation:
        if conversation_id and conversation_id in self._conversations:
            conv = self._conversations[conversation_id]
            conv.updated_at = datetime.now(timezone.utc)
            return conv

        cid = conversation_id or str(uuid.uuid4())
        conv = Conversation(id=cid)
        self._conversations[cid] = conv
        # Evict oldest if over limit
        if len(self._conversations) > self._max:
            oldest_id = min(self._conversations, key=lambda k: self._conversations[k].updated_at)
            del self._conversations[oldest_id]
        return conv

    def add_message(
        self, conversation_id: str, role: str, content: str, metadata: dict | None = None
    ) -> ChatMessage:
        conv = self._conversations.get(conversation_id)
        if conv is None:
            raise KeyError(f"Conversation '{conversation_id}' not found")
        msg = ChatMessage(role=role, content=content, metadata=metadata or {})
        conv.messages.append(msg)
        conv.updated_at = datetime.now(timezone.utc)
        return msg

    def get_history(self, conversation_id: str, limit: int = 50) -> list[ChatMessage]:
        conv = self._conversations.get(conversation_id)
        if conv is None:
            return []
        return conv.messages[-limit:]

    def list_conversations(self, limit: int = 20) -> list[Conversation]:
        sorted_convs = sorted(
            self._conversations.values(), key=lambda c: c.updated_at, reverse=True
        )
        return sorted_convs[:limit]


# Module-level singleton
_conversation_store = ConversationStore()


# ---------------------------------------------------------------------------
# Capability mapping
# ---------------------------------------------------------------------------

CAPABILITY_CATEGORIES: dict[str, dict[str, str]] = {
    "chat": {
        "display_name": "智能对话",
        "description": "自然语言交互，自动路由到对应能力",
        "icon": "MessageOutlined",
    },
    "workflow": {
        "display_name": "工作流编排",
        "description": "可视化拖拽编排工作流",
        "icon": "ApartmentOutlined",
    },
    "knowledge": {
        "display_name": "知识库",
        "description": "文档摄取、语义检索、多源RAG",
        "icon": "BookOutlined",
    },
    "skills": {
        "display_name": "技能管理",
        "description": "浏览和管理已注册的技能",
        "icon": "AppstoreOutlined",
    },
    "terminal": {
        "display_name": "智能终端",
        "description": "交互式终端会话和命令执行",
        "icon": "CodeOutlined",
    },
    "computer_use": {
        "display_name": "Computer Use",
        "description": "UI自动化操作和截屏识别",
        "icon": "DesktopOutlined",
    },
    "evolution": {
        "display_name": "自进化",
        "description": "经验积累、避坑预警、路径优化",
        "icon": "RiseOutlined",
    },
    "settings": {
        "display_name": "系统设置",
        "description": "配置LLM、技能、知识库连接",
        "icon": "SettingOutlined",
    },
}


# ---------------------------------------------------------------------------
# Request / Response models
# ---------------------------------------------------------------------------


class ChatRequest(BaseModel):
    message: str
    conversation_id: str | None = None
    sources: list[str] | None = None
    skill_name: str | None = None


class ChatResponse(BaseModel):
    conversation_id: str
    message: str
    matched_skill: str | None = None
    routing_method: str | None = None
    confidence: float | None = None
    task_id: str | None = None
    status: str = "completed"


class CapabilityInfo(BaseModel):
    name: str
    display_name: str
    description: str
    icon: str
    enabled: bool
    skill_count: int


class CapabilitiesResponse(BaseModel):
    capabilities: list[CapabilityInfo]


# ---------------------------------------------------------------------------
# Helper: resolve agent + skill for a chat request
# ---------------------------------------------------------------------------


async def _resolve_for_chat(
    request: ChatRequest, req: Request
) -> tuple[Any, Any, str | None, str | None, float | None]:
    """Resolve agent and skill for a chat request.

    Returns (agent, skill, matched_skill_name, routing_method, confidence).
    """
    pool = req.app.state.agent_pool
    skill_registry = req.app.state.skill_registry
    intent_router: IntentRouter = req.app.state.intent_router

    matched_skill_name: str | None = None
    routing_method: str | None = None
    confidence: float | None = None

    if request.skill_name:
        # Use specified skill directly
        try:
            skill = skill_registry.get(request.skill_name)
        except Exception:
            raise HTTPException(
                status_code=404,
                detail=f"Skill '{request.skill_name}' not found",
            )
        matched_skill_name = request.skill_name
        routing_method = "direct"
        confidence = 1.0
        agent = pool.get_agent(request.skill_name)
        if agent is None:
            agent = await pool.create_agent_from_skill(request.skill_name)
        return agent, skill, matched_skill_name, routing_method, confidence

    # Use IntentRouter
    all_skills = skill_registry.list_skills()
    if not all_skills:
        raise HTTPException(
            status_code=400,
            detail="No skills available. Please register skills first.",
        )
    try:
        routing_result = await intent_router.route(
            {"query": request.message, "sources": request.sources}, all_skills
        )
        matched_skill_name = routing_result.matched_skill
        routing_method = routing_result.method
        confidence = routing_result.confidence
        skill = skill_registry.get(matched_skill_name)
        agent = pool.get_agent(matched_skill_name)
        if agent is None:
            agent = await pool.create_agent_from_skill(matched_skill_name)
    except (ValueError, RuntimeError) as e:
        raise HTTPException(status_code=400, detail=str(e))

    return agent, skill, matched_skill_name, routing_method, confidence


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------


@router.post("/portal/chat", response_model=ChatResponse)
async def chat(request: ChatRequest, req: Request, _auth: None = Depends(_verify_api_key)):
    """Send a chat message and get a response with intent routing."""
    agent, skill, matched_skill, routing_method, confidence = await _resolve_for_chat(
        request, req
    )

    # Create or reuse conversation
    conv = _conversation_store.get_or_create(request.conversation_id)
    _conversation_store.add_message(conv.id, "user", request.message)

    # Build task and execute
    task = TaskMessage(
        task_id=str(uuid.uuid4()),
        agent_name=agent.name,
        task_type=agent.agent_type,
        priority=0,
        input_data={"query": request.message, "sources": request.sources},
        callback_url=None,
        created_at=datetime.now(timezone.utc),
    )

    task_result = await agent.execute(task)

    # Extract response text
    if task_result.output_data:
        if isinstance(task_result.output_data, dict):
            response_text = task_result.output_data.get("result") or task_result.output_data.get(
                "output"
            ) or json.dumps(task_result.output_data, ensure_ascii=False)
        else:
            response_text = str(task_result.output_data)
    elif task_result.error_message:
        response_text = task_result.error_message
    else:
        response_text = ""

    _conversation_store.add_message(conv.id, "assistant", response_text)

    return ChatResponse(
        conversation_id=conv.id,
        message=response_text,
        matched_skill=matched_skill,
        routing_method=routing_method,
        confidence=confidence,
        task_id=task.task_id,
        status="completed",
    )


@router.post("/portal/chat/stream")
async def chat_stream(request: ChatRequest, req: Request, _auth: None = Depends(_verify_api_key)):
    """Stream chat responses via SSE."""
    from sse_starlette.sse import EventSourceResponse

    agent, skill, matched_skill, routing_method, confidence = await _resolve_for_chat(
        request, req
    )

    # Create or reuse conversation
    conv = _conversation_store.get_or_create(request.conversation_id)
    _conversation_store.add_message(conv.id, "user", request.message)

    async def event_generator():
        react_config = agent.get_react_config()
        react_engine = ReActEngine(
            llm_gateway=req.app.state.llm_gateway,
            max_steps=react_config["max_steps"],
        )

        messages = [{"role": "user", "content": request.message}]
        tools = agent.get_tools()
        model = agent.get_model()
        system_prompt = agent.get_system_prompt()
        timeout_seconds = react_config["timeout_seconds"]

        # Send routing info as first event
        yield {
            "event": "routing",
            "data": json.dumps(
                {
                    "skill": matched_skill,
                    "method": routing_method,
                    "confidence": confidence,
                }
            ),
        }

        collected_output: list[str] = []
        try:
            async for event in react_engine.execute_stream(
                messages=messages,
                tools=tools,
                model=model,
                agent_name=agent.name,
                system_prompt=system_prompt,
                timeout_seconds=timeout_seconds,
            ):
                if event.event_type == "final_answer":
                    collected_output.append(
                        event.data.get("output", "")
                    )
                yield {
                    "event": event.event_type,
                    "data": json.dumps(
                        {
                            "step": event.step,
                            "data": event.data,
                            "timestamp": event.timestamp,
                        }
                    ),
                }
        except Exception as e:
            yield {
                "event": "error",
                "data": json.dumps({"error": str(e)}),
            }
            return

        # Save assistant response to conversation
        response_text = "".join(collected_output) if collected_output else ""
        if response_text:
            _conversation_store.add_message(conv.id, "assistant", response_text)

    return EventSourceResponse(event_generator())


@router.get("/portal/capabilities", response_model=CapabilitiesResponse)
async def get_capabilities(req: Request, _auth: None = Depends(_verify_api_key)):
    """List all available capabilities with their status."""
    skill_registry = req.app.state.skill_registry
    all_skills = skill_registry.list_skills()

    # Build a map of capability tag -> skill count
    cap_skill_counts: dict[str, int] = {}
    for skill in all_skills:
        for cap in skill.capabilities:
            cap_skill_counts[cap.tag] = cap_skill_counts.get(cap.tag, 0) + 1
        # Also count the skill itself toward "skills" category
        cap_skill_counts["skills"] = cap_skill_counts.get("skills", 0) + 1

    capabilities: list[CapabilityInfo] = []
    for cat_name, cat_info in CAPABILITY_CATEGORIES.items():
        skill_count = cap_skill_counts.get(cat_name, 0)
        capabilities.append(
            CapabilityInfo(
                name=cat_name,
                display_name=cat_info["display_name"],
                description=cat_info["description"],
                icon=cat_info["icon"],
                enabled=True,
                skill_count=skill_count,
            )
        )

    return CapabilitiesResponse(capabilities=capabilities)


@router.get("/portal/conversations")
async def list_conversations(limit: int = 20, _auth: None = Depends(_verify_api_key)):
    """List recent conversations."""
    convs = _conversation_store.list_conversations(limit=limit)
    return [
        {
            "id": c.id,
            "created_at": c.created_at.isoformat(),
            "updated_at": c.updated_at.isoformat(),
            "message_count": len(c.messages),
        }
        for c in convs
    ]


@router.get("/portal/conversations/{conversation_id}")
async def get_conversation(conversation_id: str, limit: int = 50, _auth: None = Depends(_verify_api_key)):
    """Get conversation history."""
    history = _conversation_store.get_history(conversation_id, limit=limit)
    if not history and conversation_id not in _conversation_store._conversations:
        raise HTTPException(status_code=404, detail=f"Conversation '{conversation_id}' not found")
    return [
        {
            "role": m.role,
            "content": m.content,
            "timestamp": m.timestamp.isoformat(),
            "metadata": m.metadata,
        }
        for m in history
    ]


@router.websocket("/portal/ws")
async def portal_websocket(websocket: WebSocket):
    """Real-time chat WebSocket endpoint."""
    # Authentication
    configured_api_key: str | None = None
    if hasattr(websocket.app.state, "server_config") and websocket.app.state.server_config:
        configured_api_key = websocket.app.state.server_config.api_key
    if configured_api_key is None and hasattr(websocket.app.state, "api_key"):
        configured_api_key = websocket.app.state.api_key

    # Check api_key query param
    if configured_api_key:
        provided = websocket.query_params.get("api_key")
        if not hmac.compare_digest((provided or "").encode(), configured_api_key.encode()):
            await websocket.close(code=4001, reason="Invalid or missing api_key")
            return

    await websocket.accept()

    # Create conversation
    conv = _conversation_store.get_or_create()
    await websocket.send_json({"type": "connected", "conversation_id": conv.id})

    try:
        while True:
            try:
                raw = await asyncio.wait_for(websocket.receive_text(), timeout=120.0)
            except asyncio.TimeoutError:
                await websocket.close(code=1000, reason="Heartbeat timeout")
                return

            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                continue

            msg_type = msg.get("type")

            if msg_type == "cancel":
                await websocket.send_json(
                    {"type": "result", "data": {"status": "cancelled"}}
                )
                return

            if msg_type != "chat":
                continue

            message_text = msg.get("message", "")
            sources = msg.get("sources")

            if not message_text:
                continue

            # Unified routing via CostAwareRouter (handles Layer 0/1/2)
            pool = websocket.app.state.agent_pool
            skill_registry = websocket.app.state.skill_registry
            llm_gateway = websocket.app.state.llm_gateway
            intent_router: IntentRouter = websocket.app.state.intent_router
            cost_aware_router = websocket.app.state.cost_aware_router

            all_skills = skill_registry.list_skills()
            if not all_skills:
                await websocket.send_json(
                    {"type": "error", "data": {"message": "No skills available"}}
                )
                continue

            # Get default tools for CostAwareRouter routing (only if default skill exists)
            default_tools = []
            default_system_prompt = None
            default_agent = pool.get_agent("default")
            if default_agent is not None:
                default_tools = default_agent.get_tools()
                default_system_prompt = default_agent.get_system_prompt()
            else:
                # Fallback to first available skill's tools
                for skill in all_skills:
                    agent = pool.get_agent(skill.name)
                    if agent is not None:
                        default_tools = agent.get_tools()
                        default_system_prompt = agent.get_system_prompt()
                        break

            # Route via CostAwareRouter (Layer 0/1/2)
            routing_result = await cost_aware_router.route(
                content=message_text,
                skill_registry=skill_registry,
                intent_router=intent_router,
                default_tools=default_tools,
                default_system_prompt=default_system_prompt,
                default_model="default",
                default_agent_name="default",
                session_id=conv.id,
                transparency="SILENT",
            )

            await websocket.send_json({
                "type": "routing",
                "skill": routing_result.agent_name or "default",
                "method": routing_result.match_method or "intent",
                "confidence": routing_result.match_confidence,
            })

            # Execute based on routing method
            if routing_result.match_method in ("greeting", "chat_mode"):
                # Zero-cost path: direct LLM call, no ReAct loop
                response = await llm_gateway.chat(
                    messages=[{"role": "user", "content": message_text}],
                    model="default",
                    agent_name="default",
                    task_type="chat",
                )
                await websocket.send_json({
                    "type": "result",
                    "data": {"status": "completed", "content": response.content},
                })
                continue

            # General path: agent execution
            agent_name = routing_result.agent_name or "default"
            agent = pool.get_agent(agent_name)
            if agent is None:
                agent = await pool.create_agent_from_skill(agent_name)

            # Execute via ReAct stream
            react_config = agent.get_react_config()
            react_engine = ReActEngine(
                llm_gateway=llm_gateway,
                max_steps=react_config["max_steps"],
            )

            messages = [{"role": "user", "content": message_text}]
            tools = agent.get_tools()
            model = agent.get_model()
            system_prompt = agent.get_system_prompt()
            timeout_seconds = react_config["timeout_seconds"]
            logger.info(
                f"[portal] agent='{agent_name}' tools={len(tools)} "
                f"[{', '.join(t.name for t in tools)}] model={model}"
            )
            collected_output: list[str] = []
            try:
                async for event in react_engine.execute_stream(
                    messages=messages,
                    tools=tools,
                    model=model,
                    agent_name=agent.name,
                    system_prompt=system_prompt,
                    timeout_seconds=timeout_seconds,
                ):
                    if event.event_type == "final_answer":
                        collected_output.append(event.data.get("output", ""))
                    await websocket.send_json(
                        {
                            "type": "step",
                            "data": {
                                "event_type": event.event_type,
                                "step": event.step,
                                "data": event.data,
                                "timestamp": event.timestamp,
                            },
                        }
                    )
            except Exception as e:
                await websocket.send_json(
                    {"type": "error", "data": {"message": str(e)}}
                )
                continue

            response_text = "".join(collected_output) if collected_output else ""
            if response_text:
                _conversation_store.add_message(conv.id, "assistant", response_text)

            await websocket.send_json(
                {"type": "result", "data": {"message": response_text}}
            )

    except WebSocketDisconnect:
        logger.debug(f"Portal WebSocket disconnected for conversation {conv.id}")
    except Exception as e:
        logger.error(f"Portal WebSocket error: {e}")
        try:
            await websocket.send_json(
                {"type": "error", "data": {"message": str(e)}}
            )
        except Exception:
            pass
