"""
Work Assistant Agent Server

FastAPI 기반 HTTP 서버로 에이전트 기능 제공
"""

import os
import json
import asyncio
from typing import Optional, Dict, Any, List
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from sse_starlette.sse import EventSourceResponse
from dotenv import load_dotenv
import uvicorn

from .agent import WorkAssistantAgent, AgentConfig, AgentContext, create_agent


# 환경 변수 로드
load_dotenv()


# =============================================================================
# Pydantic 모델
# =============================================================================

class ChatRequest(BaseModel):
    """채팅 요청"""
    message: str = Field(..., description="사용자 메시지")
    tenant_id: str = Field(default="", description="테넌트 ID")
    user_uid: str = Field(default="", description="사용자 UID")
    user_email: str = Field(default="", description="사용자 이메일")
    user_name: str = Field(default="", description="사용자 이름")
    conversation_id: Optional[str] = Field(default=None, description="대화 ID (세션 유지용)")
    stream: bool = Field(default=True, description="스트리밍 여부")
    metadata: Dict[str, Any] = Field(default_factory=dict, description="추가 메타데이터")


class ChatResponse(BaseModel):
    """채팅 응답"""
    response: str = Field(..., description="에이전트 응답")
    conversation_id: str = Field(..., description="대화 ID")
    tool_calls: List[Dict[str, Any]] = Field(default_factory=list, description="호출된 도구 목록")


class MCPServerRequest(BaseModel):
    """MCP 서버 추가 요청"""
    name: str = Field(..., description="서버 이름")
    command: str = Field(..., description="실행 명령어")
    args: List[str] = Field(default_factory=list, description="명령어 인자")
    env: Dict[str, str] = Field(default_factory=dict, description="환경 변수")


class HealthResponse(BaseModel):
    """헬스 체크 응답"""
    status: str = "ok"
    version: str = "0.1.0"
    mcp_servers: List[str] = Field(default_factory=list)
    tools_count: int = 0


# =============================================================================
# 글로벌 상태
# =============================================================================

# 에이전트 인스턴스
agent: Optional[WorkAssistantAgent] = None

# 대화 세션 저장소 (메모리)
conversations: Dict[str, AgentContext] = {}


# =============================================================================
# 애플리케이션 설정
# =============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """애플리케이션 라이프사이클 관리"""
    global agent
    
    # 시작 시
    print("[Server] 에이전트 초기화 중...")
    
    config = AgentConfig(
        model_provider=os.getenv("MODEL_PROVIDER", "openai"),
        model_name=os.getenv("MODEL_NAME", "gpt-4o"),
        temperature=float(os.getenv("TEMPERATURE", "0")),
        mcp_config_path=os.getenv("MCP_CONFIG_PATH", "mcp_config.json"),
        streaming=True
    )
    
    agent = WorkAssistantAgent(config)
    
    try:
        await agent.initialize()
        print(f"[Server] 에이전트 초기화 완료 (도구: {len(agent.tools)}개)")
    except Exception as e:
        print(f"[Server] 에이전트 초기화 경고: {e}")
        print("[Server] MCP 없이 기본 모드로 시작합니다.")
        # MCP 설정 없이도 기본 동작 가능
        if agent.llm is None:
            agent.llm = agent._create_llm()
        agent._initialized = True
    
    yield
    
    # 종료 시
    print("[Server] 에이전트 정리 중...")
    if agent:
        await agent.cleanup()
    print("[Server] 서버 종료")


app = FastAPI(
    title="Work Assistant Agent",
    description="LangChain ReAct Agent with MCP Support",
    version="0.1.0",
    lifespan=lifespan
)

# CORS 설정
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# =============================================================================
# 엔드포인트
# =============================================================================

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """헬스 체크"""
    mcp_servers = []
    tools_count = 0
    
    if agent and agent.mcp_client:
        mcp_servers = list(agent.mcp_client.sessions.keys())
        tools_count = len(agent.tools)
    
    return HealthResponse(
        status="ok",
        version="0.1.0",
        mcp_servers=mcp_servers,
        tools_count=tools_count
    )


@app.get("/tools")
async def list_tools():
    """사용 가능한 도구 목록"""
    if not agent:
        raise HTTPException(status_code=503, detail="에이전트가 초기화되지 않았습니다.")
    
    tools = []
    for tool in agent.tools:
        tools.append({
            "name": tool.name,
            "description": tool.description,
            "schema": tool.args_schema.model_json_schema() if tool.args_schema else {}
        })
    
    return {"tools": tools, "count": len(tools)}


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """채팅 (비스트리밍)"""
    if not agent:
        raise HTTPException(status_code=503, detail="에이전트가 초기화되지 않았습니다.")
    
    # 대화 컨텍스트 가져오기 또는 생성
    conversation_id = request.conversation_id or _generate_conversation_id()
    
    if conversation_id not in conversations:
        conversations[conversation_id] = AgentContext(
            tenant_id=request.tenant_id,
            user_uid=request.user_uid,
            user_email=request.user_email,
            user_name=request.user_name,
            metadata=request.metadata
        )
    
    context = conversations[conversation_id]
    
    # 에이전트 실행
    try:
        response = await agent.run(request.message, context, stream=False)
        
        return ChatResponse(
            response=response,
            conversation_id=conversation_id,
            tool_calls=[]
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/chat/stream")
async def chat_stream(request: ChatRequest):
    """채팅 (스트리밍)"""
    if not agent:
        raise HTTPException(status_code=503, detail="에이전트가 초기화되지 않았습니다.")
    
    # 대화 컨텍스트 가져오기 또는 생성
    conversation_id = request.conversation_id or _generate_conversation_id()
    
    if conversation_id not in conversations:
        conversations[conversation_id] = AgentContext(
            tenant_id=request.tenant_id,
            user_uid=request.user_uid,
            user_email=request.user_email,
            user_name=request.user_name,
            metadata=request.metadata
        )
    
    context = conversations[conversation_id]
    
    async def event_generator():
        """SSE 이벤트 생성기"""
        try:
            # 대화 ID 먼저 전송
            yield {
                "event": "metadata",
                "data": json.dumps({"conversation_id": conversation_id})
            }
            
            # 스트리밍 응답
            stream = await agent.run(request.message, context, stream=True)
            
            async for chunk in stream:
                yield {
                    "event": "message",
                    "data": chunk
                }
            
        except Exception as e:
            yield {
                "event": "error",
                "data": json.dumps({"error": str(e)})
            }
    
    return EventSourceResponse(event_generator())


@app.post("/mcp/servers")
async def add_mcp_server(request: MCPServerRequest):
    """MCP 서버 추가"""
    if not agent:
        raise HTTPException(status_code=503, detail="에이전트가 초기화되지 않았습니다.")
    
    try:
        await agent.add_mcp_server(
            name=request.name,
            command=request.command,
            args=request.args,
            env=request.env
        )
        
        return {
            "status": "success",
            "message": f"MCP 서버 '{request.name}' 추가됨",
            "tools_count": len(agent.tools)
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/mcp/servers")
async def list_mcp_servers():
    """등록된 MCP 서버 목록"""
    if not agent or not agent.mcp_client:
        return {"servers": [], "connected": []}
    
    servers = []
    for name, config in agent.mcp_client.servers.items():
        servers.append({
            "name": name,
            "command": config.command,
            "args": config.args,
            "enabled": config.enabled
        })
    
    connected = list(agent.mcp_client.sessions.keys())
    
    return {"servers": servers, "connected": connected}


@app.delete("/conversations/{conversation_id}")
async def delete_conversation(conversation_id: str):
    """대화 세션 삭제"""
    if conversation_id in conversations:
        del conversations[conversation_id]
        return {"status": "deleted", "conversation_id": conversation_id}
    else:
        raise HTTPException(status_code=404, detail="대화를 찾을 수 없습니다.")


@app.get("/conversations")
async def list_conversations():
    """대화 세션 목록"""
    return {
        "conversations": [
            {
                "id": conv_id,
                "tenant_id": ctx.tenant_id,
                "user_email": ctx.user_email,
                "message_count": len(ctx.conversation_history)
            }
            for conv_id, ctx in conversations.items()
        ]
    }


# =============================================================================
# 유틸리티
# =============================================================================

def _generate_conversation_id() -> str:
    """대화 ID 생성"""
    import uuid
    return str(uuid.uuid4())


# =============================================================================
# 메인
# =============================================================================

def main():
    """서버 실행"""
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", "8000"))
    
    print(f"[Server] Starting on {host}:{port}")
    uvicorn.run(
        "work_assistant_agent.server:app",
        host=host,
        port=port,
        reload=os.getenv("DEBUG", "false").lower() == "true"
    )


if __name__ == "__main__":
    main()

