from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager
from sqlalchemy.orm import Session
from api import users, chatrooms, admin, ai_agents, weather, chat_assistant
import models, schemas
from database import get_db, engine, Base
from typing import Dict, List

# Create tables
Base.metadata.create_all(bind=engine)

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理，只在 worker 进程中运行一次"""
    try:
        from init_builtin_tools import init_builtin_tools
        print("正在加载内置工具...")
        init_builtin_tools()
    except Exception as e:
        print(f"加载内置工具失败: {e}")
    yield  # 应用运行中

app = FastAPI(title="AI Customer Service System API", version="1.0.0", lifespan=lifespan)

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

# 注册路由
app.include_router(users.router, prefix="/api/users", tags=["Users"])
app.include_router(chatrooms.router, prefix="/api/chatrooms", tags=["Chatrooms"])
app.include_router(admin.router, prefix="/api/admin", tags=["Admin"])
app.include_router(ai_agents.router, prefix="/api/ai-agents", tags=["AI Agents"])
app.include_router(weather.router, prefix="/api", tags=["Weather"])
app.include_router(chat_assistant.router, prefix="/api/chat-assistant", tags=["Chat Assistant"])

@app.get("/")
def read_root():
    return {"message": "Welcome to AI Customer Service System API"}

# --- WebSocket Chat Manager ---
class ConnectionManager:
    def __init__(self):
        # room_id -> list of websockets
        self.active_connections: Dict[int, List[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, room_id: int):
        await websocket.accept()
        if room_id not in self.active_connections:
            self.active_connections[room_id] = []
        self.active_connections[room_id].append(websocket)

    def disconnect(self, websocket: WebSocket, room_id: int):
        if room_id in self.active_connections:
            self.active_connections[room_id].remove(websocket)
            if not self.active_connections[room_id]:
                del self.active_connections[room_id]

    async def broadcast_to_room(self, message: dict, room_id: int):
        if room_id in self.active_connections:
            for connection in self.active_connections[room_id]:
                await connection.send_json(message)

manager = ConnectionManager()

@app.websocket("/ws/chat/{room_id}/{user_id}")
async def websocket_endpoint(websocket: WebSocket, room_id: int, user_id: int, db: Session = Depends(get_db)):
    await manager.connect(websocket, room_id)
    try:
        while True:
            data = await websocket.receive_text()
            
            # 尝试解析JSON格式的消息（用于表情图片）
            message_type = "text"
            message_content = data
            
            try:
                import json
                msg_data = json.loads(data)
                if isinstance(msg_data, dict) and "type" in msg_data and "content" in msg_data:
                    message_type = msg_data["type"]
                    message_content = msg_data["content"]
                    data = message_content
            except (json.JSONDecodeError, ValueError):
                # 不是JSON格式，按普通文本处理
                pass
            
            # 检查是否是数字员工的消息（以 [AI- 开头）
            is_employee_message = data.startswith("[AI-")
            sender_name = "Unknown"
            actual_sender_id = user_id
            
            if is_employee_message:
                # 解析数字员工消息格式: [AI-员工名] 消息内容
                import re
                match = re.match(r'\[AI-(.+?)\]\s*(.*)', data)
                if match:
                    employee_name = match.group(1)
                    actual_content = match.group(2)
                    
                    # 查找数字员工
                    employee = db.query(models.DigitalEmployee).filter(
                        models.DigitalEmployee.name == employee_name
                    ).first()
                    
                    if employee:
                        # 使用负数 ID 表示数字员工
                        actual_sender_id = -employee.id
                        sender_name = employee.name
                        data = actual_content
                    else:
                        # 如果找不到数字员工，使用原始消息
                        user = db.query(models.User).filter(models.User.id == user_id).first()
                        sender_name = user.username if user else "Unknown"
                else:
                    user = db.query(models.User).filter(models.User.id == user_id).first()
                    sender_name = user.username if user else "Unknown"
            else:
                # 普通用户消息
                user = db.query(models.User).filter(models.User.id == user_id).first()
                sender_name = user.username if user else "Unknown"
            
            # Save message to DB
            new_msg = models.Message(
                room_id=room_id, 
                sender_id=actual_sender_id, 
                content=data,
                message_type=message_type
            )
            db.add(new_msg)
            db.commit()
            db.refresh(new_msg)
            
            # Broadcast to everyone in the room
            msg_payload = {
                "id": new_msg.id,
                "room_id": room_id,
                "sender_id": actual_sender_id,
                "sender_name": sender_name,
                "content": data,
                "message_type": message_type,
                "timestamp": new_msg.timestamp.isoformat(),
                "is_sensitive": False,
                "is_employee": actual_sender_id < 0
            }
            await manager.broadcast_to_room(msg_payload, room_id)
            
    except WebSocketDisconnect:
        manager.disconnect(websocket, room_id)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)