#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Time    :   2024/04/10 11:07:17
@Author  :   ChenHao
@Description  :   智谱api
@Contact :   jerrychen1990@gmail.com
'''


import inspect
import itertools
import json
import time
from typing import Any, Iterable, List, Tuple

import numpy as np

from agit.common import LLMResp, ToolCall, ToolDesc, Usage

from loguru import logger
import os
from typing import List
from cachetools import LRUCache, cached
from loguru import logger

import zhipuai
from agit import AGIT_LOG_HOME
from snippets import jdumps, retry, batch_process
from zhipuai import ZhipuAI


fmt = "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> <level>{level: <8}</level>|  - <level>[{extra[request_id]}]{message}</level>"

logger.add(os.path.join(AGIT_LOG_HOME, "zhipuai_bk.log"), retention="365 days", rotation=" 00:00", level="DEBUG", format=fmt,
           filter=lambda r: r["function"] == "call_llm_api_v4")

default_logger = logger


def check_api_key(api_key):
    if api_key is None:
        api_key = os.environ.get("ZHIPUAI_API_KEY", None)
    if not api_key:
        raise ValueError("未设置api_key且未设置ZHIPUAI_API_KEY环境变量")
    zhipuai.api_key = api_key


@cached(LRUCache(maxsize=1000))
def get_client(api_key: str) -> ZhipuAI:
    logger.debug(f"new zhipuai client with {api_key=}")
    client = ZhipuAI(api_key=api_key or os.environ.get("ZHIPU_API_KEY"))
    return client


def support_system(model: str):
    return model.startswith("chatglm3") or model in ["glm-4", "glm-3-turbo", "glm-4v"]


def _is_multimodal(model: str):
    return model in ["glm-4v"]


def build_messages(messages: List[dict], model: str, system: str):
    if system:
        if support_system(model):
            messages = [dict(role="system", content=system)] + messages
        else:
            logger.warning(f"{model} not support system message, system argument will not work")

    detail_msgs = []
    for idx, item in enumerate(messages):
        detail_msgs.append(f"[{idx+1}].{item['role']}:{item['content']}")
    logger.debug("\n"+"\n".join(detail_msgs))

    return messages


def resp2generator(api_resp: Iterable, first_chunk, request_id: str, usage: Usage):
    def gen():
        acc = []
        for chunk in itertools.chain([first_chunk], api_resp):
            choices = chunk.choices
            if choices:
                choice = choices[0]
                if choice.delta.content:
                    delta_content = choice.delta.content
                    yield delta_content
                    acc.append(delta_content)
            _usage = chunk.usage
            if _usage:
                usage.__dict__.update(_usage.model_dump())

        resp_msg = "".join(acc).strip()
        with logger.contextualize(request_id=request_id):
            logger.debug(f"model generate answer:{resp_msg}")

    return gen()

def build_tool_calls(tool_calls) -> List[ToolCall]:
    # logger.debug(f"tool_calls: {tool_calls}")
    if tool_calls is None:
        return []
    rs = []
    for tool_call in tool_calls:
        parameters = json.loads(tool_call.function.arguments)   

        tmp = ToolCall(tool_call_id= tool_call.id, name=tool_call.function.name, parameters=parameters)
        rs.append(tmp)
    # logger.debug(f"tool_calls: {rs}")
    return rs

def _convert_tool_desc(tools: List[ToolDesc]) -> List[dict]:
    resp = []
    for tool in tools:

        properties = {p.name: dict(type=p.type, description=p.description) for p in tool.parameters}
        required = [p.name for p in tool.parameters if p.required]
        parameters = dict(type="object", properties=properties, required=required)
        tool_desc = dict(type="function", function=dict(name=tool.name, description=tool.description, parameters=parameters))
        resp.append(tool_desc)

    return resp

def _extract_tool_calls(api_resp)->Tuple[List[ToolCall], Any]:
    tool_calls = []
    for chunk in api_resp:
        choices = chunk.choices
        if choices:
            choice = choices[0]
            tool_calls = choice.delta.tool_calls
            tool_calls = build_tool_calls(tool_calls)
        break            
    return tool_calls, chunk


def call_llm(messages: List[str], model: str, request_id: str,
             system=None, stream=True, api_key=None,
             do_search=False, search_query=None, tools: List[ToolDesc] = [],
             temperature=0.7, top_p=0.95, max_tokens: int = None, do_sample=True,
             **kwargs) -> LLMResp:

    st = time.time()
    with logger.contextualize(request_id=request_id):
        client = get_client(api_key=api_key or os.environ.get("ZHIPU_API_KEY"))
        messages = build_messages(messages=messages, model=model, system=system)

        tools = _convert_tool_desc(tools=tools)

        if not do_search and model == "glm-4":
            close_search_tool = {'type': 'web_search', 'web_search': {'enable': False, 'search_query': search_query}}
            tools.append(close_search_tool)

        params = dict(temperature=temperature, top_p=top_p, max_tokens=max_tokens, do_sample=do_sample, **kwargs)
        # 处理temperature=0的特殊情况
        if params.get("temperature") == 0:
            params.pop("temperature")
            params["do_sample"] = False

        # 过滤无效的参数
        valid_keys = dict(inspect.signature(client.chat.completions.create).parameters).keys()
        params = {k: v for k, v in params.items() if k in valid_keys}

        api_inputs = dict(model=model,
                          messages=messages,
                          tools=tools,
                          stream=stream,
                          **params)
        logger.debug(f"api_inputs:\n{jdumps(api_inputs)}")
        details = dict(api_inputs=api_inputs)
        api_resp = client.chat.completions.create(
            **api_inputs
        )
        if stream:
            usage = Usage(total_tokens=None, completion_tokens=None, prompt_tokens=None)
            tool_calls,first_chunk = _extract_tool_calls(api_resp=api_resp)
     
            content = resp2generator(api_resp, first_chunk, request_id, usage)

        else:
            logger.debug(f"api_outputs:\n{jdumps(api_resp)}")

            content = api_resp.choices[0].message.content
            tool_calls = build_tool_calls(api_resp.choices[0].message.tool_calls)
            usage = Usage.model_validate(api_resp.usage.model_dump())
            details.update(api_resp=api_resp)

        # total_cost = time.time() - st
        # total_speed = usage.total_tokens/total_cost
        # perf = Perf(first_token_latency=first_token_latency, total_cost=total_cost, decode_speed=decode_speed, encode_speed=encode_speed, total_speed=total_speed)
        resp = LLMResp(content=content, details=details, usage=usage, perf=None, tool_calls=tool_calls)
        return resp


def _call_embedding_single(text:str, api_key:str=None, model:str="embedding-2",
                           norm=True, retry_num=2, wait_time=(1,2), **kwargs)->List[float]:
    client = get_client(api_key=api_key or os.environ.get("ZHIPU_API_KEY"))

    def zhipu_embedding_attempt():
        resp = client.embeddings.create(
            model=model,  # 填写需要调用的模型名称
            input=text
        )
        embedding = resp.data[0].embedding
        if norm:
            embedding = embedding / np.linalg.norm(embedding, 2)
            embedding = embedding.tolist()
        return embedding

    if retry_num:
        attempt = retry(retry_num=retry_num, wait_time=wait_time)(zhipu_embedding_attempt)
    return attempt()

def call_embedding(text: str | List[str], api_key: str=None, model: str="embedding-2",  norm=True,
                   batch_size=8, retry_num=2, wait_time=(1, 2)) -> List[float] | List[List[float]]:
    if isinstance(text, str):
        return _call_embedding_single(text=text, model=model, api_key=api_key, norm=norm, retry_num=retry_num, wait_time=wait_time)
    else:
        batch_embd_func = batch_process(work_num=batch_size, return_list=True)(_call_embedding_single)
        return batch_embd_func(data=text, model=model, api_key=api_key, norm=norm, retry_num=retry_num, wait_time=wait_time)
