from pydantic import BaseModel, field_validator, model_validator, Field
from typing import List, Optional, Generic, TypeVar
from fastapi import FastAPI
import uvicorn
import logging


from qwenvl import model
from qwenvl import processor
from qwen_vl_utils import process_vision_info

logger = logging.getLogger()
logger.setLevel(logging.INFO)
handler1 = logging.StreamHandler()
handler2 = logging.FileHandler(filename='../log/llmserver.log')
formatter = logging.Formatter(
    "%(asctime)s - %(module)s - %(funcName)s - line:%(lineno)d - %(levelname)s - %(message)s"
)
handler1.setFormatter(formatter)
handler2.setFormatter(formatter)
logger.addHandler(handler1)  # 将日志输出至屏幕
logger.addHandler(handler2)  # 将日志输出至文件


app = FastAPI()

T = TypeVar("T")

class APIRequest(BaseModel):
    imageData : str
    text  : str

class APIResponse(BaseModel, Generic[T]):
    success: bool
    data: Optional[List[T]]
    msg: Optional[List[str]]

@app.post("/llm/detect")
@app.post("/llm/detect/")
async def detect(item: APIRequest):
    # illegal 为 0 代表没有违章
    # illegal 为 1 代表有违章
    response = {
        "sucess": "OK", 
        "data": {"illegal" : 0}, 
        "msg": ""
    }
    
    # 提示词
    prompt_text  = item.text
    base64_image = item.imageData
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": f"data:image;base64,{base64_image}",
                },
                {"type": "text", "text": prompt_text},
            ],
        }
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    inputs = inputs.to("cuda")

    # Inference: Generation of the output
    generated_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids_trimmed = [
        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    print(output_text)
    # 大模型检测后如果有违章
    # response["data"]["illegal"] = 1
    if "yes" in output_text[0].lower():
        response["data"]["illegal"] = 1
    return response

if __name__ == "__main__":
    uvicorn.run('server:app', host="0.0.0.0", port=18000)