import sys
import importlib
from pydantic import BaseModel
from flask import Flask, Response, request
from flask_cors import CORS
import json

sys.path.append("../")
from qwen_agent.planning.plan_executor import PlanExecutor
from qwen_agent.planning.plan_continue_executor import PlanContinueExecutor

prompt_lan = "CN"
llm_name = "qwen-plus"
max_ref_token = 4000
model_server = "http://10.10.0.10:7907/v1"
api_key = ""
server_host = "127.0.0.1"

app = Flask(__name__)
CORS(app)

rspHeaders = {
    "Cache-Control": "no-cache",
    "Connection": "keep-alive",
    "Transfer-Encoding": "chunked",
}

if model_server.startswith("http"):
    source = "local"
elif model_server.startswith("dashscope"):
    source = "dashscope"

if llm_name.startswith("gpt"):
    module = "qwen_agent.llm.gpt"
    llm = importlib.import_module(module).GPT(llm_name)
elif llm_name.startswith("Qwen") or llm_name.startswith("qwen"):
    module = "qwen_agent.llm.qwen"
    llm = importlib.import_module(module).Qwen(
        llm_name, model_server=model_server, api_key=api_key
    )
else:
    raise NotImplementedError

planContinueExecutor = PlanContinueExecutor(enable_critic=False, llm=llm, stream=True)
planExecutor = PlanExecutor(enable_critic=False, llm=llm, stream=True)


@app.route("/subscribe/<question>", methods=["GET", "POST"])
def subscribe(question: str):
    return Response(
        call_with_stream(question),
        mimetype="text/event-stream",
        headers=rspHeaders,
    )


class ClarificationRequest(BaseModel):
    data: str


@app.route("/clarification/", methods=["GET", "POST"])
def clarification():
    print("clarification request: ", request)
    data = request.get_json("data")
    print("clarification: ", data)
    return Response(
        call_with_stream(json.dumps(data), True),
        mimetype="text/event-stream",
        headers=rspHeaders,
    )


def call_with_stream(question, isClarification=False):
    if isClarification:
        executor = planContinueExecutor
    else:
        executor = planExecutor
    try:
        for rsp in executor.run(question, []):
            yield f"data:{rsp}"
            yield "data: [DONE]"
    except Exception as error:
        print("error: ", error)
        yield "data: [DONE]"


if __name__ == "__main__":
    # app.run(host="0.0.0.0", port=20020, debug=False, processes=10, threaded=False)
    app.run(host="0.0.0.0", port=20020, debug=False)