import re
import os
import json
import tabulate as tb
import requests
import numpy as np
from prompt import doc_generate_prompt, daoci_prompt, material_retrieval_prompt

CHARSET = 'UTF-8'
BAIDU_ACCESS_TOKEN = os.getenv('BAIDU_ACCESS_TOKEN')
GOOGLE_SEARCH_API = 'https://content-customsearch.googleapis.com/customsearch/v1'
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
GOOGLE_CX = os.getenv('GOOGLE_CX')

BAIDU_APIs = {
    'ErnieBot': {
        'url': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions',
        'params': {
            'messages': {
                'type': list,
                'optional': False,
                'default': None,
                'description': '聊天上下文信息。'
            },
            'temperature': {
                'type': float,
                'optional': True,
                'default': 0.95,
                'description': '说明：'
                               '（1）较高的数值会使输出更加随机，而较低的数值会使其更加集中和确定。'
                               '（2）默认0.95，范围 (0, 1.0]，不能为0。'
                               '（3）建议该参数和top_p只设置1个。'
                               '（4）建议top_p和temperature不要同时更改。。'
            },
            'top_p': {
                'type': float,
                'optional': True,
                'default': 0.8,
                'description': '说明：'
                               '（1）影响输出文本的多样性，取值越大，生成文本的多样性越强。'
                               '（2）默认0.8，取值范围 [0, 1.0]。'
                               '（3）建议该参数和temperature只设置1个。'
                               '（4）建议top_p和temperature不要同时更改。'
            },
            'penalty_score': {
                'name': '',
                'type': float,
                'optional': True,
                'default': 1.0,
                'description': '通过对已生成的token增加惩罚，减少重复生成的现象。'
                               '说明：'
                               '（1）值越大表示惩罚越大。'
                               '（2）默认1.0，取值范围：[1.0, 2.0]。'
            },
        }
    },
    'ErnieBot-turbo': {
        'url': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant',
        'params': {
            'messages': {
                'type': list,
                'optional': False,
                'default': None,
                'description': '聊天上下文信息。'
            },
            'user_id': {
                'type': str,
                'optional': True,
                'default': None,
                'description': '表示最终用户的唯一标识符，可以监视和检测滥用行为，防止接口恶意调用。'
            },
        },
        'description': ''
    },
    'Bloomz-7B': {
        'url': 'https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/bloomz_7b1',
        'params': {
            'messages': {
                'type': list,
                'optional': False,
                'default': None,
                'description': '聊天上下文信息。'
            },
            'user_id': {
                'type': str,
                'optional': True,
                'default': None,
                'description': '表示最终用户的唯一标识符，可以监视和检测滥用行为，防止接口恶意调用。'
            },
        },
        'description': ''
    },
    'title_generate': {
        'url': 'https://aip.baidubce.com/rpc/2.0/nlp/v1/titlepredictor',
        'params': {
            'doc': {
                'type': str,
                'optional': False,
                'default': None,
                'description': '输入文本，200字到1500字之间，默认GBK编码；若须输入UTF-8编码文本，请添加URL参数charset=UTF-8'
            }
        },
        'description': ''
    },
    'summary': {
        'url': 'https://aip.baidubce.com/rpc/2.0/nlp/v1/news_summary',
        'params': {
            'content': {
                'type': str,
                'optional': False,
                'default': None,
                'description': '(Required) 字符串（限3000字符数以内） 字符串仅支持GBK编码，长度需小于3000字符数（即6000字节），'
                               '请输入前确认字符数没有超限，若字符数超长会返回错误。'
                               '正文中如果包含段落信息，请使用"\n"分隔，段落信息算法中有重要的作用，请尽量保留'
            },
            'max_summary_len': {
                'type': int,
                'optional': False,
                'default': None,
                'description': '(Required) 此数值将作为摘要结果的最大长度。'
                               '例如：原文长度1000字，本参数设置为150，则摘要结果的最大长度是150字；推荐最优区间：200-500字'
            },
            'title': {
                'type': int,
                'optional': True,
                'default': None,
                'description': '字符串（限200字符数） 字符串仅支持GBK编码，长度需小于200字符数（即400字节），'
                               '请输入前确认字符数没有超限，若字符数超长会返回错误。'
                               '标题在算法中具有重要的作用，若文章确无标题，不必输入此参数'
            }
        },
        'description': ''
    },
}


def get(url, params):
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/67.0.3396.87 "
                      "Safari/537.36",
    }
    new_url = url + '?'
    for k, v in params.items():
        new_url += f'&{k}={v}'

    response = requests.get(new_url, headers=headers).json()
    return response


def post(url, body):
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }
    payload = json.dumps(body)
    new_url = url + f'?access_token={BAIDU_ACCESS_TOKEN}&charset={CHARSET}'
    response = requests.request("POST", new_url, headers=headers, data=payload).json()
    return response


def doc_generate(date, site, persons, records, doc_type, model, temperature, penalty_score):
    url = BAIDU_APIs[model]['url']
    prompt = doc_generate_prompt(date, site, persons, records, doc_type)
    body = {
        'messages': [
            {
                'role': 'user',
                'content': prompt
            }
        ]
    }

    if model in ['ErnieBot']:
        body['temperature'] = temperature
        body['penalty_score'] = penalty_score

    response = post(url, body)

    if 'result' in response.keys():
        return True, response['result']
    else:
        return False, response['error_msg']


def summary(content, maxlen, title=None):
    url = BAIDU_APIs['summary']['url']

    body = {
        'content': content,
        'max_summary_len': maxlen
    }

    if title:
        body['title'] = title

    response = post(url, body)

    if 'summary' in response.keys():
        return True, response['summary']
    else:
        return False, response['error_msg']


def standardization(data):
    mu = np.mean(data, axis=0)
    sigma = np.std(data, axis=0)
    return (data - mu) / sigma


def softmax(x):
    return np.exp(x) / np.sum(np.exp(x), axis=0)


def title_generate(doc):
    url = BAIDU_APIs['title_generate']['url']

    body = {
        'doc': doc
    }

    response = post(url, body)

    if 'reference_titles' in response.keys():
        titles = []
        scores = []
        for item in response['reference_titles']:
            titles.append(item['title'])
            scores.append(item['score'])
        scores = softmax(standardization(scores))
        return True, dict(zip(titles, scores))
    else:
        return False, response['error_msg']


def material_retrieval(doc, kw_num, m_num):
    model = 'ErnieBot-turbo'
    prompt = material_retrieval_prompt(doc)
    print(prompt)
    body = {
        'messages': [
            {
                'role': 'user',
                'content': prompt
            }
        ]
    }

    try:
        url = BAIDU_APIs[model]['url']
        res = post(url, body)['result']
        print(res)
        j = json.loads(re.findall(r'```json[\s\S]*```', res)[0][7:-3])
        keywords = []
        print(j)
        if type(j)==dict:
            for k, v in j.items():
                keywords += v
        elif type(j)==list:
            keywords += j
        else:
            return False, '模型输出格式错误!'
        if len(keywords) > kw_num:
            keywords = keywords[:kw_num]
    except Exception:
        return False, '提取失败'

    result = {}
    for keyword in keywords:
        params = {
            'filter': 1,
            'siteSearch': '*.gov.cn',
            'siteSearchFilter': 'i',
            'q': keyword,
            'num': m_num,
            'key': GOOGLE_API_KEY,
            'cx': GOOGLE_CX
        }

        response = get(GOOGLE_SEARCH_API, params)

        if 'error' in response.keys():
            return False, response['error']['message']
        else:
            items = response['items']
            new_items = [
                {
                    '标题': item['title'],
                    '链接': item['link'],
                    '摘要': item['snippet']
                } for item in items
            ]
            result[keyword] = new_items

    # 写markdown

    ms = ''

    for keyword in result.keys():
        ms += f'## 关键词： {keyword}\n'
        for i, m in enumerate(result[keyword]):
            ms += f"- 素材{i+1}\n- - 标题：{m['标题']}\n- - 链接：<{m['链接']}>\n- - 摘要：{m['摘要']}\n"
        ms += '\n'
    out = f"{ms}"
    print(out)
    return True, ('、'.join(keywords), out)


def daoci_generate(name, gender, age, political, dead_time, dead_reason, experience, model, temperature, penalty_score):
    url = BAIDU_APIs[model]['url']
    prompt = daoci_prompt(name, gender, age, political, dead_time, dead_reason, experience)
    body = {
        'messages': [
            {
                'role': 'user',
                'content': prompt
            }
        ]
    }

    if model in ['ErnieBot']:
        body['temperature'] = temperature
        body['penalty_score'] = penalty_score

    response = post(url, body)

    if 'result' in response.keys():
        return True, response['result']
    else:
        return False, response['error_msg']


if __name__ == '__main__':
    prompt = """
    从以下文字中提取有关具体党政机关、地方政策的关键词，用于在互联网上检索相关信息进一步学习，请确保提取的关键词是具体的政策、思想等可以进一步学习的知识性事物：

“紧紧围绕高质量发展这个全面建设社会主义现代化国家的首要任务，以强化理论学习指导发展实践，以深化调查研究推动解决发展难题，把学习和调研落实到完成党的二十大部署的各项任务中去，以推动高质量发展的新成效检验主题教育成果。”

　　在学习贯彻习近平新时代中国特色社会主义思想主题教育工作会议上，习近平总书记明确把“推动高质量发展”作为这次主题教育的一项重点措施并提出系列要求。这次主题教育总要求中明确要“重实践、建新功”，具体目标中有“实干担当促进发展”，工作安排中有“推动发展”，指向是一致的，就是要教育实践两手抓、两促进，在改造主观世界同时改造客观世界，全力推动高质量发展取得新成效。

　　发展是我们党执政兴国的第一要务，是解决我国一切问题的基础和关键。推动高质量发展，是以习近平同志为核心的党中央深入分析我国发展新的历史条件和阶段、全面认识和把握我国现代化建设实践历程以及各国现代化建设一般规律作出的重大决策部署。党的十八大以来，我们党对经济发展阶段性特征的认识不断深化，就高质量发展作出一系列重大判断，指出我国经济发展已由高速增长阶段转向高质量发展阶段，强调要坚持以推动高质量发展为主题，提出构建新发展格局是贯彻新发展理念、推动高质量发展的重大举措，明确实现高质量发展是中国式现代化的本质要求、是全面建设社会主义现代化国家的首要任务。习近平总书记指出，“实现高质量发展，是保持经济社会持续健康发展的必然要求，是适应我国社会主要矛盾变化和全面建设社会主义现代化国家的必然要求”；“高质量发展，就是能够很好满足人民日益增长的美好生活需要的发展，是体现新发展理念的发展，是创新成为第一动力、协调成为内生特点、绿色成为普遍形态、开放成为必由之路、共享成为根本目的的发展”；“高质量发展是‘十四五’乃至更长时期我国经济社会发展的主题，关系我国社会主义现代化建设全局”……总书记围绕为什么要推动高质量发展、什么是高质量发展、怎样推动高质量发展等问题作出一系列重要论述，为新时代推动高质量发展提供了行动纲领和根本遵循。

　　科学理论引领伟大实践。党的十八大以来，以习近平同志为核心的党中央统筹推进“五位一体”总体布局，协调推进“四个全面”战略布局，推动我国经济实力、科技实力、综合国力跃上新的大台阶，沿着高质量发展之路阔步前行。新时代十年，国内生产总值增加近70万亿元，在高基数基础上实现了中高速增长、迈向高质量发展；科技创新成果丰硕，一些关键核心技术实现突破，进入创新型国家行列；区域协调发展机制不断健全，区域重大战略稳步实施；绿水青山就是金山银山的理念深入人心，生态环境保护发生历史性、转折性、全局性变化；实行更加积极主动的开放战略，全方位高水平开放型经济加快形成；建成世界上规模最大的教育体系、社会保障体系、医疗卫生体系，人民生活全方位改善。特别是我们历史性地解决了绝对贫困问题，如期全面建成小康社会，实现第一个百年奋斗目标，我国发展站在新的更高历史起点上。

　　现在，我们党带领全国各族人民迈上了全面建设社会主义现代化国家、全面推进中华民族伟大复兴的新征程。同时，世界百年未有之大变局加速演进，不确定、难预料因素增多，国内改革发展稳定面临不少深层次矛盾躲不开、绕不过，各种风险挑战、困难问题比以往更加严峻复杂。面对新的战略机遇、新的战略任务、新的战略阶段、新的战略要求、新的战略环境，我们必须牢牢把握高质量发展这个首要任务，不断夯实全面建成社会主义现代化强国的物质技术基础。

　　习近平总书记强调，“这次主题教育，要教育引导广大党员、干部学思想、见行动”。全力推动高质量发展，是这次主题教育的重要着力点。各地各部门开展主题教育的成效如何，最终还是要用结果来评判，要落实到推动高质量发展的行动上。要看到，当前我国发展不平衡不充分问题仍然突出，实现高质量发展还有许多卡点瓶颈。广大党员、干部要胸怀“国之大者”，坚持学用结合、以学促干，紧紧围绕新时代新征程党的中心任务，围绕推动高质量发展的实践要求，增强推动高质量发展本领、服务群众本领、防范化解风险本领，以“时时放心不下”的责任感、积极担当作为的精气神为党和人民履好职、尽好责，以新气象新作为推动高质量发展取得新成效，依靠顽强斗争打开事业发展新天地。

输出格式：以json格式输出一个包含关键字字符串的列表。

关键字限定个数：不得超过5个。
    """
    model = 'ErnieBot-turbo'

    body = {
        'messages': [
            {
                'role': 'user',
                'content': prompt
            }
        ]
    }

    if model in ['ErnieBot']:
        body['temperature'] = 0.95
        body['penalty_score'] = 1.0

    url = BAIDU_APIs[model]['url']

    res = post(url, body)['result']

    print(res)