from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from kafka import KafkaProducer, KafkaAdminClient
from kafka.admin import NewTopic
import uuid
import json
import mysql.connector
from Config.Config import KAFKA_BOOTSTRAP_SERVERS, KAFKA_TOPIC, MYSQL_CONFIG

# 定义请求体模型
class TaskRequest(BaseModel):
    prompt: str  # 用户输入的提示词

# 初始化 FastAPI 应用
app = FastAPI()

# 创建 Kafka 生产者
producer = KafkaProducer(
    bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS,
    value_serializer=lambda v: json.dumps(v).encode('utf-8')  # 将消息序列化为 JSON
)

# 检查并创建主题
def ensure_topic_exists():
    try:
        # 创建 Kafka AdminClient
        admin_client = KafkaAdminClient(bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS)

        # 获取所有主题
        existing_topics = admin_client.list_topics()

        # 如果主题不存在，则创建
        if KAFKA_TOPIC not in existing_topics:
            topic = NewTopic(
                name=KAFKA_TOPIC,
                num_partitions=1,  # 分区数
                replication_factor=1  # 副本数
            )
            admin_client.create_topics([topic])
            print(f"Topic '{KAFKA_TOPIC}' created successfully.")
        else:
            print(f"Topic '{KAFKA_TOPIC}' already exists.")

        # 关闭 AdminClient
        admin_client.close()
    except Exception as e:
        print(f"Failed to ensure topic exists: {str(e)}")

# 定义 Lifespan 事件
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 应用启动时检查并创建主题
    ensure_topic_exists()
    yield
    # 应用关闭时清理资源（可选）
    print("Application shutdown.")

# 写入 MySQL 日志
def log_to_mysql(task_id: str, keyword: str):
    try:
        # 连接 MySQL
        connection = mysql.connector.connect(**MYSQL_CONFIG)
        cursor = connection.cursor()

        # 插入日志
        query = """
        INSERT INTO t_gen_tasks (task_id, keyword, status)
        VALUES (%s, %s, %s)
        """
        cursor.execute(query, (task_id, keyword, 'pending'))
        connection.commit()

        # 关闭连接
        cursor.close()
        connection.close()
    except Exception as e:
        print(f"Failed to log to MySQL: {str(e)}")

# 定义接口
@app.post("/create-task")
def create_task(task_request: TaskRequest):
    try:
        # 生成唯一的 task_id
        task_id = str(uuid.uuid4())

        # 构造任务信息
        task_data = {
            "task_id": task_id,
            "prompt": task_request.prompt
        }

        # 将任务信息写入 Kafka
        producer.send(KAFKA_TOPIC, value=task_data)
        producer.flush()  # 确保消息发送完成

        # 将任务信息写入 MySQL
        log_to_mysql(task_id, task_request.prompt)

        # 返回 task_id
        return {"task_id": task_id}
    except Exception as e:
        # 处理异常
        raise HTTPException(status_code=500, detail=f"Failed to create task: {str(e)}")

# 启动应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)