from typing import List, Optional, Literal, Union, Dict, Any

from fastapi import APIRouter, Depends
import numpy
from pydantic import BaseModel, Field, TypeAdapter

from embedding_sever.api.resp_bean import RespBean, success

router = APIRouter()


# ============== 请求/响应模型定义 ==============

class FaceKwargs(BaseModel):
    """人脸类型扩展字段"""
    person_id: Optional[str] = Field(None, description="人物id")
    person_name: Optional[str] = Field(None, description="人物姓名")


class SportShotKwargs(BaseModel):
    """体育镜头类型扩展字段"""
    match_name: Optional[str] = Field(None, description="比赛名称")
    person_name: Optional[str] = Field(None, description="人名")
    video_desc: Optional[str] = Field(None, description="视频描述")
    shot_cls: Optional[str] = Field(None, description="镜头分类")
    match_id: Optional[str] = Field(None, description="比赛id")
    program_id: Optional[str] = Field(None, description="节目id")


# 定义类型映射，用于根据type自动转换kwargs
KwargsType = Union[FaceKwargs, SportShotKwargs]


class PutDataRequest(BaseModel):
    """写入数据请求体"""
    id: str = Field(..., description="唯一键id")
    embedding: List[float] = Field(..., description="向量")
    embedding_version: str = Field(..., description="向量版本 小写数字下划线组成")
    kwargs: Optional[Dict[str, Any]] = Field(None, description="扩展字段，根据type不同字段不同")

    def get_typed_kwargs(self, type: str) -> Optional[KwargsType]:
        """根据type自动转换kwargs为具体类型"""
        if self.kwargs is None:
            return None
        type_map = {
            'face': FaceKwargs,
            'sport_shot': SportShotKwargs
        }
        kwargs_class = type_map.get(type)
        if kwargs_class:
            return kwargs_class(**self.kwargs)
        return None


class DelByIdRequest(BaseModel):
    """删除数据请求体"""
    ids: List[str] = Field(..., description="id 集合")
    embedding_version: str = Field(..., description="向量版本 小写数字下划线组成")


class FilterItem(BaseModel):
    """过滤条件项"""
    name: str = Field(..., description="字段名")
    value: List = Field(..., description="字段值")
    opt: Optional[Literal["eq", "neq", "lt", "gt", "lte", "gte", "like", "in"]] = Field(
        None, description="操作类型: eq相等, neq不等, lt小于, gt大于, lte小于等于, gte大于等于, like模糊匹配, in内"
    )


class SearchRequest(BaseModel):
    """检索数据请求体"""
    embedding: List[float] = Field(..., description="向量")
    embedding_version: str = Field(..., description="向量版本 小写数字下划线组成")
    topk: int = Field(..., description="top k")
    filters: Optional[List[FilterItem]] = Field(None, description="过滤字段")


class SearchResultItem(BaseModel):
    """检索结果单项"""
    id: str = Field(..., description="id")
    kwargs: dict = Field(..., description="扩展字段")
    score: float = Field(..., description="分值")


class SearchResponseData(BaseModel):
    """检索响应数据"""
    type: str = Field(..., description="type")
    data_list: List[SearchResultItem] = Field(..., description="topk结果列表")


# ============== 接口定义 ==============
from embedding_sever.service.data_service import DataService
from embedding_sever.service import get_data_service


@router.put("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据")
@router.post("/{type}/put", response_model=RespBean, tags=["数据服务"], summary="写入数据")
async def put_data(
        type: Literal["face", "sport_shot"],
        request: PutDataRequest,
        data_service: DataService = Depends(get_data_service)
):
    """
    写入数据
    """
    request.embedding = numpy.array(request.embedding)

    embedding_version = request.embedding_version
    tb_name = f'{type}_{embedding_version}'
    upserted_id = await data_service.upsert(tb_name, request.id, request.embedding,
                                            request.get_typed_kwargs(type).model_dump())

    return success(data={'id': upserted_id})


@router.delete("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据")
@router.post("/{type}/del_by_id", response_model=RespBean, tags=["数据服务"], summary="删除数据")
async def del_by_id(
        type: Literal["face", "sport_shot"],
        request: DelByIdRequest,
        data_service: DataService = Depends(get_data_service)
):
    """
    删除数据
    
    - **type**: 数据类型 (face-人脸, sport_shot-体育镜头)
    - **ids**: id集合
    """
    embedding_version = request.embedding_version
    tb_name = f'{type}_{embedding_version}'
    await data_service.delete_by_pks(tb_name, request.ids)

    return success()


@router.post("/{type}/search", response_model=RespBean[SearchResponseData], tags=["数据服务"], summary="检索数据")
async def search_data(
        type: Literal["face", "sport_shot"],
        request: SearchRequest,
        data_service: DataService = Depends(get_data_service),
        adapter=TypeAdapter(List[SearchResultItem])
):
    """
    检索数据
    
    - **type**: 数据类型 (face-人脸, sport_shot-体育镜头)
    - **embedding**: 查询向量
    - **embedding_version**: 向量版本
    - **topk**: 返回top k条结果
    - **filters**: 过滤条件列表
    """
    #
    embedding_version = request.embedding_version
    tb_name = f'{type}_{embedding_version}'

    embedding = request.embedding
    filters = request.filters
    from_, size = 0, request.topk

    result_data = await data_service.search(tb_name, embedding, filters, from_, size)
    search_results: List[SearchResultItem] = adapter.validate_python(result_data)

    return success(data=SearchResponseData(type=type, data_list=search_results))