from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
import json
import asyncio
from typing import List, Dict
import random

app = FastAPI(title="Chat API with SSE")

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


# Simulate AI responses based on user questions
def get_ai_response(messages: List[Dict]) -> str:
    """Generate AI response based on message history"""
    if not messages:
        return "你好！我是AI助手，有什么可以帮助你的吗？"

    last_message = messages[-1]["content"].lower()

    # Simulate different responses based on keywords
    if "销售额" in last_message or "销售" in last_message:
        return """根据数据分析：

去年(2024年)销售额：2,850万元
前年(2023年)销售额：2,340万元

对比分析：
- 增长额：510万元
- 增长率：21.8%
- 主要增长来源：线上渠道增长35%，新产品线贡献40%

整体趋势良好，建议继续保持线上渠道投入。"""

    elif "个体工商户" in last_message or "小微企业" in last_message:
        return """表2.4 个体工商户及小微企业主基本情况的采集范围：

1. **基本信息采集**
   - 营业执照信息（注册号、名称、经营范围）
   - 经营者个人信息（姓名、身份证号、联系方式）
   - 经营地址及联系方式

2. **经营状况采集**
   - 从业人员数量
   - 年营业收入
   - 资产总额
   - 经营年限

3. **行业分类信息**
   - 所属行业（按国民经济行业分类）
   - 主营业务描述
   - 经营模式（线上/线下/混合）

4. **财务信息**
   - 纳税情况
   - 贷款需求
   - 政策扶持情况

此采集范围适用于年营业收入在500万元以下的个体工商户及小微企业主。"""

    elif "你好" in last_message or "hi" in last_message.lower():
        return "你好！我是AI助手，很高兴为您服务。您可以问我关于数据分析、业务咨询等问题。"

    else:
        return f"我理解您询问的是「{messages[-1]['content']}」。这是一个很好的问题。基于您的问题，我建议从以下几个方面考虑：\n\n1. 首先需要明确具体的业务需求和目标\n2. 收集相关的数据和资料进行分析\n3. 制定详细的实施方案\n4. 持续跟进和优化\n\n如果您需要更具体的帮助，请提供更多细节信息。"


async def event_generator(messages: List[Dict]):
    """Generate SSE events for streaming response"""
    try:
        # Send start event
        yield f"data: {json.dumps({'type': 'start'})}\n\n"

        # Get AI response
        response = get_ai_response(messages)

        # Stream response character by character with some delay to simulate thinking
        for i, char in enumerate(response):
            # Send content event
            yield f"data: {json.dumps({'type': 'content', 'content': char})}\n\n"

            # Add small delay for streaming effect
            # Faster delay for spaces and punctuation
            if char in [' ', '\n', '，', '。', '、', '：', '；']:
                await asyncio.sleep(0.01)
            else:
                await asyncio.sleep(0.02)

        # Send end event
        yield f"data: {json.dumps({'type': 'end'})}\n\n"

    except Exception as e:
        # Send error event
        yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"


@app.get("/chat")
async def chat(messages: str):
    """
    Chat endpoint with SSE streaming

    Args:
        messages: JSON string of message history

    Returns:
        StreamingResponse with SSE events
    """
    try:
        # Parse messages
        messages_list = json.loads(messages)

        return StreamingResponse(
            event_generator(messages_list),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no"
            }
        )
    except json.JSONDecodeError:
        return {"error": "Invalid JSON format for messages"}
    except Exception as e:
        return {"error": str(e)}


@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": "Chat API with SSE",
        "endpoints": {
            "/chat": "GET - Chat with SSE streaming (requires 'messages' query param)",
        }
    }


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