#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
Date: 2021/11/23 20:31
Desc: 股票基本信息
"""
import json
import warnings
from io import BytesIO

import pandas as pd
import requests
from tqdm import tqdm


def stock_info_sz_name_code(indicator: str = "A股列表") -> pd.DataFrame:
    """
    深圳证券交易所-股票列表
    http://www.szse.cn/market/product/stock/list/index.html
    :param indicator: choice of {"A股列表", "B股列表", "CDR列表", "AB股列表"}
    :type indicator: str
    :return: 指定 indicator 的数据
    :rtype: pandas.DataFrame
    """
    url = "http://www.szse.cn/api/report/ShowReport"
    indicator_map = {"A股列表": "tab1", "B股列表": "tab2", "CDR列表": "tab3", "AB股列表": "tab4"}
    params = {
        "SHOWTYPE": "xlsx",
        "CATALOGID": "1110",
        "TABKEY": indicator_map[indicator],
        "random": "0.6935816432433362",
    }
    r = requests.get(url, params=params)
    with warnings.catch_warnings(record=True):
        warnings.simplefilter("always")
        temp_df = pd.read_excel(BytesIO(r.content))
    if len(temp_df) > 10:
        if indicator == "A股列表":
            temp_df["A股代码"] = (
                temp_df["A股代码"]
                .astype(str)
                .str.split(".", expand=True)
                .iloc[:, 0]
                .str.zfill(6)
                .str.replace("000nan", "")
            )
            temp_df = temp_df[
                [
                    "板块",
                    "A股代码",
                    "A股简称",
                    "A股上市日期",
                    "A股总股本",
                    "A股流通股本",
                    "所属行业",
                ]
            ]
        elif indicator == "B股列表":
            temp_df["B股代码"] = (
                temp_df["B股代码"]
                .astype(str)
                .str.split(".", expand=True)
                .iloc[:, 0]
                .str.zfill(6)
                .str.replace("000nan", "")
            )
            temp_df = temp_df[
                [
                    "板块",
                    "B股代码",
                    "B股简称",
                    "B股上市日期",
                    "B股总股本",
                    "B股流通股本",
                    "所属行业",
                ]
            ]
        elif indicator == "AB股列表":
            temp_df["A股代码"] = (
                temp_df["A股代码"]
                .astype(str)
                .str.split(".", expand=True)
                .iloc[:, 0]
                .str.zfill(6)
                .str.replace("000nan", "")
            )
            temp_df["B股代码"] = (
                temp_df["B股代码"]
                .astype(str)
                .str.split(".", expand=True)
                .iloc[:, 0]
                .str.zfill(6)
                .str.replace("000nan", "")
            )
            temp_df = temp_df[
                [
                    "板块",
                    "A股代码",
                    "A股简称",
                    "A股上市日期",
                    "B股代码",
                    "B股简称",
                    "B股上市日期",
                    "所属行业",
                ]
            ]
        return temp_df
    else:
        return temp_df


def stock_info_sh_name_code(indicator: str = "主板A股") -> pd.DataFrame:
    """
    上海证券交易所-股票列表
    http://www.sse.com.cn/assortment/stock/list/share/
    :param indicator: choice of {"主板A股": "1", "主板B股": "2", "科创板": "8"}
    :type indicator: str
    :return: 指定 indicator 的数据
    :rtype: pandas.DataFrame
    """
    indicator_map = {"主板A股": "1", "主板B股": "2", "科创板": "8"}
    url = "http://query.sse.com.cn/security/stock/getStockListData.do"
    headers = {
        "Host": "query.sse.com.cn",
        "Pragma": "no-cache",
        "Referer": "http://www.sse.com.cn/assortment/stock/list/share/",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36",
    }
    params = {
        "jsonCallBack": "jsonpCallback66942",
        "isPagination": "true",
        "stockCode": "",
        "csrcCode": "",
        "areaName": "",
        "stockType": indicator_map[indicator],
        "pageHelp.cacheSize": "1",
        "pageHelp.beginPage": "1",
        "pageHelp.pageSize": "2000",
        "pageHelp.pageNo": "1",
        "pageHelp.endPage": "11",
        "_": "1589881387934",
    }
    r = requests.get(url, params=params, headers=headers)
    text_data = r.text
    json_data = json.loads(text_data[text_data.find("{") : -1])
    temp_df = pd.DataFrame(json_data["result"])
    temp_df.columns = [
        '-',
        '公司简称',
        '-',
        '-',
        '-',
        '-',
        '-',
        '-',
        '-',
        '简称',
        '代码',
        '-',
        '-',
        '公司代码',
        '-',
        '上市日期',
    ]
    temp_df = temp_df[[
        '公司代码',
        '公司简称',
        '代码',
        '简称',
        '上市日期',
    ]]
    temp_df['上市日期'] = pd.to_datetime(temp_df['上市日期']).dt.date
    return temp_df


def stock_info_bj_name_code() -> pd.DataFrame:
    """
    北京证券交易所-股票列表
    http://www.bse.cn/nq/listedcompany.html
    :return: 股票列表
    :rtype: pandas.DataFrame
    """
    url = "http://www.bse.cn/nqxxController/nqxxCnzq.do"
    payload = {
        "page": "0",
        "typejb": "T",
        "xxfcbj[]": "2",
        "xxzqdm": "",
        "sortfield": "xxzqdm",
        "sorttype": "asc",
    }
    r = requests.post(url, data=payload)
    data_text = r.text
    data_json = json.loads(data_text[data_text.find("[") : -1])
    total_page = data_json[0]["totalPages"]
    big_df = pd.DataFrame()
    for page in tqdm(range(total_page)):
        payload.update({"page": page})
        r = requests.post(url, data=payload)
        data_text = r.text
        data_json = json.loads(data_text[data_text.find("[") : -1])
        temp_df = data_json[0]["content"]
        big_df = big_df.append(temp_df, ignore_index=True)
    big_df.columns = [
        "上市日期",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "流通股本",
        "-",
        "-",
        "-",
        "-",
        "-",
        "所属行业",
        "-",
        "-",
        "-",
        "-",
        "报告日期",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "地区",
        "-",
        "-",
        "-",
        "-",
        "-",
        "券商",
        "总股本",
        "-",
        "证券代码",
        "-",
        "证券简称",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
        "-",
    ]
    big_df = big_df[
        [
            "证券代码",
            "证券简称",
            "总股本",
            "流通股本",
            "上市日期",
            "所属行业",
            "地区",
            "报告日期",
        ]
    ]
    big_df['报告日期'] = pd.to_datetime(big_df['报告日期']).dt.date
    big_df['上市日期'] = pd.to_datetime(big_df['上市日期']).dt.date
    return big_df


def stock_info_sh_delist(indicator: str = "终止上市公司") -> pd.DataFrame:
    """
    上海证券交易所-暂停上市公司-终止上市公司
    http://www.sse.com.cn/assortment/stock/list/firstissue/
    :param indicator: choice of {"终止上市公司": "5", "暂停上市公司": "4"}
    :type indicator: str
    :return: 暂停上市公司 or 终止上市公司 的数据
    :rtype: pandas.DataFrame
    """
    indicator_map = {"终止上市公司": "5", "暂停上市公司": "4"}
    url = "http://query.sse.com.cn/security/stock/getStockListData2.do"
    headers = {
        "Host": "query.sse.com.cn",
        "Pragma": "no-cache",
        "Referer": "http://www.sse.com.cn/assortment/stock/list/share/",
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/81.0.4044.138 Safari/537.36",
    }
    params = {
        "jsonCallBack": "jsonpCallback66942",
        "isPagination": "true",
        "stockCode": "",
        "csrcCode": "",
        "areaName": "",
        "stockType": indicator_map[indicator],
        "pageHelp.cacheSize": "1",
        "pageHelp.beginPage": "1",
        "pageHelp.pageSize": "2000",
        "pageHelp.pageNo": "1",
        "pageHelp.endPage": "11",
        "_": "1589881387934",
    }
    r = requests.get(url, params=params, headers=headers)
    text_data = r.text
    json_data = json.loads(text_data[text_data.find("{") : -1])
    temp_df = pd.DataFrame(json_data["result"])
    temp_df.columns = [
        '-',
        '-',
        '终止上市后股份转让代码',
        '-',
        '终止上市后股份转让主办券商',
        '终止上市后股份转让副主办券商',
        '终止上市日期',
        '-',
        '原公司简称',
        '原公司代码',
        '-',
        '-',
        '-',
        '-',
        '上市日期',
    ]
    temp_df = temp_df[[
        '原公司代码',
        '原公司简称',
        '上市日期',
        '终止上市日期',
        '终止上市后股份转让代码',
        '终止上市后股份转让主办券商',
        '终止上市后股份转让副主办券商',
    ]]
    temp_df['上市日期'] = pd.to_datetime(temp_df['上市日期']).dt.date
    temp_df['终止上市日期'] = pd.to_datetime(temp_df['终止上市日期']).dt.date
    return temp_df


def stock_info_sz_delist(indicator: str = "暂停上市公司") -> pd.DataFrame:
    """
    深证证券交易所-暂停上市公司-终止上市公司
    http://www.szse.cn/market/stock/suspend/index.html
    :param indicator: choice of {"暂停上市公司", "终止上市公司"}
    :type indicator: str
    :return: 暂停上市公司 or 终止上市公司 的数据
    :rtype: pandas.DataFrame
    """
    indicator_map = {"暂停上市公司": "tab1", "终止上市公司": "tab2"}
    url = "http://www.szse.cn/api/report/ShowReport"
    params = {
        "SHOWTYPE": "xlsx",
        "CATALOGID": "1793_ssgs",
        "TABKEY": indicator_map[indicator],
        "random": "0.6935816432433362",
    }
    r = requests.get(url, params=params)
    with warnings.catch_warnings(record=True):
        warnings.simplefilter("always")
        temp_df = pd.read_excel(BytesIO(r.content))
        temp_df["证券代码"] = temp_df["证券代码"].astype("str").str.zfill(6)
        return temp_df


def stock_info_sz_change_name(indicator: str = "全称变更") -> pd.DataFrame:
    """
    深证证券交易所-更名公司
    http://www.szse.cn/market/companys/changename/index.html
    :param indicator: choice of {"全称变更": "tab1", "简称变更": "tab2"}
    :type indicator: str
    :return: 全称变更 or 简称变更 的数据
    :rtype: pandas.DataFrame
    """
    indicator_map = {"全称变更": "tab1", "简称变更": "tab2"}
    url = "http://www.szse.cn/api/report/ShowReport"
    params = {
        "SHOWTYPE": "xlsx",
        "CATALOGID": "SSGSGMXX",
        "TABKEY": indicator_map[indicator],
        "random": "0.6935816432433362",
    }
    r = requests.get(url, params=params)
    with warnings.catch_warnings(record=True):
        warnings.simplefilter("always")
        temp_df = pd.read_excel(BytesIO(r.content))
        temp_df["证券代码"] = temp_df["证券代码"].astype("str").str.zfill(6)
        return temp_df


def stock_info_change_name(stock: str = "688588") -> pd.DataFrame:
    """
    新浪财经-股票曾用名
    http://vip.stock.finance.sina.com.cn/corp/go.php/vCI_CorpInfo/stockid/300378.phtml
    :param stock: 股票代码
    :type stock: str
    :return: 股票曾用名列表
    :rtype: list
    """
    url = f"http://vip.stock.finance.sina.com.cn/corp/go.php/vCI_CorpInfo/stockid/{stock}.phtml"
    r = requests.get(url)
    temp_df = pd.read_html(r.text)[3].iloc[:, :2]
    temp_df.dropna(inplace=True)
    temp_df.columns = ["item", "value"]
    temp_df["item"] = temp_df["item"].str.split("：", expand=True)[0]
    try:
        name_list = temp_df[temp_df["item"] == "证券简称更名历史"].value.tolist()[0].split(" ")
        return name_list
    except:
        return pd.DataFrame()


def stock_info_a_code_name() -> pd.DataFrame:
    """
    沪深京 A 股列表
    :return: 沪深京 A 股数据
    :rtype: pandas.DataFrame
    """
    big_df = pd.DataFrame()
    stock_sh = stock_info_sh_name_code(indicator="主板A股")
    stock_sh = stock_sh[["公司代码", "公司简称"]]

    stock_sz = stock_info_sz_name_code(indicator="A股列表")
    stock_sz["A股代码"] = stock_sz["A股代码"].astype(str).str.zfill(6)
    big_df = big_df.append(stock_sz[["A股代码", "A股简称"]], ignore_index=True)
    big_df.columns = ["公司代码", "公司简称"]

    stock_kcb = stock_info_sh_name_code(indicator="科创板")
    stock_kcb = stock_kcb[["公司代码", "公司简称"]]

    stock_bse = stock_info_bj_name_code()
    stock_bse = stock_bse[["证券代码", "证券简称"]]
    stock_bse.columns = ["公司代码", "公司简称"]

    big_df = big_df.append(stock_sh, ignore_index=True)
    big_df = big_df.append(stock_kcb, ignore_index=True)
    big_df = big_df.append(stock_bse, ignore_index=True)
    big_df.columns = ["code", "name"]
    return big_df


if __name__ == "__main__":
    stock_info_sz_df = stock_info_sz_name_code(indicator="A股列表")
    print(stock_info_sz_df)

    stock_info_sz_df = stock_info_sz_name_code(indicator="B股列表")
    print(stock_info_sz_df)

    stock_info_sz_df = stock_info_sz_name_code(indicator="AB股列表")
    print(stock_info_sz_df)

    stock_info_sz_df = stock_info_sz_name_code(indicator="CDR列表")
    print(stock_info_sz_df)

    stock_info_sh_delist_df = stock_info_sh_delist(indicator="终止上市公司")
    print(stock_info_sh_delist_df)

    stock_info_sz_delist_df = stock_info_sz_delist(indicator="终止上市公司")
    print(stock_info_sz_delist_df)

    stock_info_sz_change_name_df = stock_info_sz_change_name(indicator="全称变更")
    print(stock_info_sz_change_name_df)

    stock_info_change_name_list = stock_info_change_name(stock="000503")
    print(stock_info_change_name_list)

    stock_info_a_code_name_df = stock_info_a_code_name()
    print(stock_info_a_code_name_df)

    stock_info_bj_name_code_df = stock_info_bj_name_code()
    print(stock_info_bj_name_code_df)
