#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/10 18:33
@Author  : cao
@File    : web2.py.py
@Desc    : 
"""
import json

import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import utils
import faiss
from collections import Counter
import logging
from starlette.responses import JSONResponse
import dirty
from prometheus_fastapi_instrumentator import Instrumentator, metrics
# from prometheus_client import Counter as myCounter
# from prometheus_client import REGISTRY,gc_collector, platform_collector, process_collector


logger = logging.getLogger('my_logger')
logger.setLevel(logging.INFO)
console_handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger.addHandler(console_handler)

model = utils.load_model()
index = utils.index
cate = utils.cate

class Sentence(BaseModel):
    query: str
    sessionId: str

app = FastAPI()
Instrumentator().instrument(app).expose(app)
from registry import REQUEST_COUNT

@app.post('/intention')
def intention(sentence:Sentence):
    query, session_id = sentence.query, sentence.sessionId
    logger.info(f"intention: session_id:{session_id}, query:{query}")
    if query is None or len(query)==0:
        return JSONResponse(content={"status":0, "intention": "others", "query": query, "isFaq":False}, media_type="application/json")
    is_dirty, mess = dirty.dirty_detect(query)
    if is_dirty:
        return JSONResponse(content=mess, media_type="application/json")
    embed = model.encode([query], normalize_embeddings=True)
    D, I = index.search(embed, 3)
    index_prob = dict(zip(I[0], D[0]))
    out = []
    for doc_id, prob in index_prob.items():
        ans = cate.get(str(doc_id), {})
        ans['prob'] = prob
        ans['doc_id'] = doc_id
        out.append(ans)
    out = sorted(out, key=lambda x: float(x['prob']), reverse=True)
    res = {}
    intention = [i for i in map(lambda x: x['意图类别'], out)]
    first= out[0]['prob']
    logger.info(f"intention: session_id:{session_id}, query:{query}, first scores:{first}")
    if out[0]['prob'] >0.9 or (out[0]['prob'] >0.7 and (len(set(intention)) ==1 or len(intention)== len(set(intention)))):
        res['ans'] = out[0]['FAQ答案']
        res['doc_id'] = int(out[0]['doc_id'])
        res['score'] = float(out[0]['prob'])
        res['intent'] = out[0]['cate']
        res['class'] = out[0]['class']
        res['subclass'] = out[0]['意图类别']
        res['isFaq'] = False if out[0]['cate'] !="FAQ" else True
        res['recall'] = "emb"
        res['query'] = query
        res['status'] = 2
        res['intent'] =  out[0]['cate']
    else:
        intention_cate = dict(Counter(intention))
        intention_max = sorted(intention_cate.items(), key=lambda x: x[1], reverse=True)[0]
        out = [i for i in  filter(lambda  x: x['意图类别']==intention_max[0], out)]
        res['ans'] = out[0]['FAQ答案']
        res['doc_id'] = int(out[0]['doc_id'])
        res['score'] = float(out[0]['prob'])
        res['isFaq'] = False if out[0]['cate'] !="FAQ" else True
        res['recall'] = "emb"
        res['class'] = out[0]['class']
        res['subclass'] = out[0]['意图类别']
        res['query'] = query
        if out[0]['prob'] < 0.7:
            res['intent'] = 'others'
            REQUEST_COUNT.inc()
            res['status'] = 4
        else:
            res['intent'] = out[0]['cate']
            res['status'] = 3
    logger.info(f"intention: session_id:{session_id}: ans:{res}")
    return JSONResponse(content=res, media_type="application/json")


if __name__ == '__main__':
    uvicorn.run("web:app", host='0.0.0.0', port=50072, workers=4)