from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import InferenceClient
import uvicorn
from paho.mqtt import client as mqtt_client
import uuid
import time



app = FastAPI()
origins = [
    "*",
    "http://localhost:3000",
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

inferenceClient = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
historyMap = {}

broker_host = "broker.emqx.io"
broker_port = 1883
broker_topic_prefix = "/diopsoft/mixtral-fastapi"
# mqttClient = None
mqttClient = mqtt_client.Client(mqtt_client.CallbackAPIVersion.VERSION1, str(uuid.uuid1()))
def connectMqttBroker():
    try:
        res = mqttClient.connect(broker_host, broker_port, 60)
        print('MQTT connection status : ' + str(res))
        mqttClient.loop_start()
    except Exception as Argument:
        # logging.getLogger().exception("Failed to connect to MQTT Broker")
        # logging.getLogger().error(e)
        traceback.print_exc()
def on_connect(client, userdata, flags, rc):
        if rc == 0:
            print("Connected to MQTT Broker!")
        else:
            raise Exception("Failed to connect to MQTT Broker, return code " + rc)
mqttClient.on_connect = on_connect
# connectMqttBroker()


class Item(BaseModel):
    is_mock: bool = False
    token: str = None
    mqttTopic: str = None
    prompt: str
    history: list = None
    # system_prompt: str = "Soyez un assistant utile, en français de préférence"
    system_prompt: str = "Soyez un assistant utile"
    temperature: float = 0.7
    max_new_tokens: int = 1048
    top_p: float = 0.9
    repetition_penalty: float = 1.1

def format_prompt(message, history):
    prompt = "<s>"
    if history:
        for hist in history:
            prompt += f"[INST] {hist['user_prompt']} [/INST]"
            prompt += f" {hist['bot_response']}</s> "
    prompt += f"[INST] {message} [/INST]"
    return prompt

def generate(item: Item):
    temperature = float(item.temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(item.top_p)

    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=item.max_new_tokens,
        top_p=top_p,
        repetition_penalty=item.repetition_penalty,
        do_sample=True,
        seed=42,
    )

    if item.mqttTopic and not mqttClient.is_connected():
        connectMqttBroker()

    token = item.token
    history = None
    if token:
        history = historyMap.get(token)

    formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", history)
    if not item.is_mock:
        stream = inferenceClient.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    if not item.is_mock:
        for response in stream:
            print(response.token.text)
            output += response.token.text
            if item.mqttTopic and mqttClient.is_connected():
                # result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, response.token.text)
                result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, output)
                status = result[0]
                if status == 0:
                    print("Message sent to topic")
                else:
                    print("Failed to send message to topic")
    else:
        for i in range(15):
            text = " generated text " + str(i)
            time.sleep(0.5)
            print(text)
            output += text
            if item.mqttTopic and mqttClient.is_connected():
                # result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, text)
                result = mqttClient.publish(broker_topic_prefix+'/'+item.mqttTopic, output)
                status = result[0]
                if status == 0:
                    print("Message sent to topic")
                else:
                    print("Failed to send message to topic")

    if item.mqttTopic and mqttClient.is_connected():
        mqttClient.unsubscribe(broker_topic_prefix+'/'+item.mqttTopic)

    if token:
        if not historyMap.get(token):
            historyMap[token] = []
        historyMap.get(token).append({ "user_prompt":item.prompt, "bot_response":output })

    return output

@app.post("/generate/")
async def generate_text(item: Item):
    return {"response": generate(item)}

@app.get("/test-api/")
async def testAPI():
    return "API V2 is running"