import os
import json
import time
import yaml
import requests
import functools
from environs import Env
import lark_oapi as lark
import multiprocessing as mp
from typing import List, Dict
from flask import Flask, request
from lark_oapi.adapter.flask import *
from lark_oapi.api.im.v1 import *


env = Env()
env.read_env()

app = Flask(__name__)

lark.VERIFICATION_TOKEN = env.str("LARK_TOKEN")
lark.ENCRYPT_KEY = ""
LANGLINK_APP_ID = env.str("LANGLINK_APP_ID")
LANGLINK_OUTPUT_BLOCK = env.str("LANGLINK_OUTPUT_BLOCK")
LANGLINK_ENDPOINT = env.str("LANGLINK_ENDPOINT")
LANGLINK_ACCESS_KEY = env.str("LANGLINK_ACCESS_KEY")
LANGLINK_ACCESS_SEC = env.str("LANGLINK_ACCESS_SEC")
LANGLINK_TIMEOUT = env.int("LANGLINK_TIMEOUT")

client = lark.Client.builder() \
    .app_id(env.str("LARK_APP_ID")) \
    .app_secret(env.str("LARK_APP_SECRET")) \
    .build()

@functools.lru_cache
def get_langlink_app_info(chat_id: str) -> Dict[str, str]:
    with open(os.path.join(os.path.dirname(__file__), "tenants.yaml"), "r") as stream:
        chats = yaml.safe_load(stream)

    for chat in chats:
        if chat["chat_id"] == chat_id:
            return {
                "langlink_app_id": chat["langlink_app_id"],
                "langlink_output_block": chat["langlink_output_block"],
            }
    return {
        "langlink_app_id": LANGLINK_APP_ID,
        "langlink_output_block": LANGLINK_OUTPUT_BLOCK,
    }

def get_response_from_langlink(user_id: str, app: Dict[str, str], texts: List[str]) -> str:
    start_ts = time.time()

    url = os.path.join(LANGLINK_ENDPOINT, "applications", app["langlink_app_id"], "async")
    headers = {
        "x-langlink-access-key": LANGLINK_ACCESS_KEY,
        "x-langlink-access-secret": LANGLINK_ACCESS_SEC,
        "x-langlink-user": user_id,
    }

    r = requests.post(
        url,
        headers=headers,
        params={
            "return_trace": True,
        },
        json={
            "input": texts,
        },
    )
    if r.status_code != 200:
        return f"调用 langlink 发生错误: {r.text}"

    x = r.json()
    interaction_id = x["id"]
    debug_url = os.path.join(LANGLINK_ENDPOINT, "applications", app["langlink_app_id"], "debug", interaction_id)

    while time.time() - start_ts < LANGLINK_TIMEOUT:
        r = requests.get(
            debug_url,
            headers=headers,
        )
        if r.status_code != 200:
            return f"调用 langlink 发生错误: {r.text}"
  
        x = r.json()
        for l in x["debug"]:
            if l["block"] == app["langlink_output_block"]:
                return l["output"]

        if len(x["debug"]) > 0:
            return f"配置错误，未找到 langlink app 中的输出 block"

        time.sleep(10)

    return f"请求 langlink 超时，共计 {time.time() - start_ts}s"

def get_context(message_id: str) -> List[str]:
    r = lark.im.v1.model.GetMessageRequestBuilder().message_id(message_id).build()
    x = client.im.v1.message.get(r)
    message = x.data.items[0]
    if message.parent_id:
        xs = get_context(message.parent_id)
    else:
        xs = []
    if message.msg_type == "text":
        xs.append(json.loads(message.body.content)["text"])
    return xs

def reply_message(user_id: str, message: Message) -> None:
    texts = []
    if message.parent_id:
        texts = get_context(message.parent_id)
    if message.message_type == "text":
        content = json.loads(message.content)
        text = content["text"]
        if "@_all" in text:
            # don"t handle at all
            return None
        if message.mentions:
            for m in message.mentions:
                text = text.replace(m.key, "")
        texts.append(text)
        app = get_langlink_app_info(message.chat_id)
        r = lark.im.v1.model.ReplyMessageRequestBuilder() \
            .message_id(message.message_id) \
            .request_body(lark.im.v1.model.ReplyMessageRequestBodyBuilder() \
                            .content(json.dumps({
                                "text": get_response_from_langlink(user_id, app, texts),
                            })) \
                            .reply_in_thread(True) \
                            .msg_type("text").build() \
            ).build()
    else:
        r = lark.im.v1.model.ReplyMessageRequestBuilder() \
            .message_id(message.message_id) \
            .request_body(lark.im.v1.model.ReplyMessageRequestBodyBuilder() \
                            .content('{"text": "不支持的消息类型"}') \
                            .reply_in_thread(True) \
                            .msg_type("text").build() \
            ).build()
    client.im.v1.message.reply(r)

def do_p2_im_message_receive_v1(data: P2ImMessageReceiveV1) -> None:
    mp.Process(target=reply_message, args=(data.event.sender.sender_id.user_id, data.event.message)).start()

handler = lark.EventDispatcherHandler.builder(lark.ENCRYPT_KEY, lark.VERIFICATION_TOKEN, lark.LogLevel.DEBUG) \
    .register_p2_im_message_receive_v1(do_p2_im_message_receive_v1) \
    .build()

@app.post("/event")
def event():
    resp = handler.do(parse_req())
    return parse_resp(resp)

@app.get("/ping")
def ping():
    return {"status": "success"}