import requests
import json
import time


def get_unix_time():
    return int(time.time())

def chang_xq_stocks(stock):
    stock = stock.upper()
    if stock[:2] in ['SH', 'SZ']:
        return stock
    return ('SH' + stock[:6] if stock[-4:] == 'XSHG'    #聚宽代码
            else 'SZ' + stock[:6] if stock[-4:] == 'XSHE'
            else 'SH' + stock if int(stock[0]) >= 5 # 纯股票代码
            else 'SZ' + stock
            )
    
# 雪球自选股操作类
class XqStocks(object):
    def __init__(self, **kwargs):
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 6.1; WOW64; rv:32.0) Gecko/20100101 Firefox/32.0',
            'Host': 'xueqiu.com',
            'Pragma': 'no-cache',
            'Connection': 'keep-alive',
            'Accept': '*/*',
            'Accept-Encoding': 'gzip,deflate,sdch',
            'Cache-Control': 'no-cache',
            'Referer': 'https://xueqiu.com/portfolios',
            'X-Requested-With': 'XMLHttpRequest',
            'Accept-Language': 'zh-CN,zh;q=0.5',
            'Content-Type': 'application/x-www-form-urlencoded; charset=UTF-8'
        }
        self.session = requests.Session()
        self.session.headers.update(headers)
        if 'proxy_ip' in kwargs:
            s = "%s:%s" % (kwargs['proxy_ip'], kwargs.get('proxy_port',808))
            if 'proxy_user' in kwargs and 'proxy_pwd' in kwargs:
                s = '%s:%s@%s' % (kwargs['proxy_user'], kwargs['proxy_pwd'], s)
            self.session.proxies = {
                "http": "http://" + s,
                "https": "https://" + s,
            }
        self.account_config = kwargs
        self.config = {
            'login_api': 'https://xueqiu.com/user/login',
            'stocks': 'https://xueqiu.com/v4/stock/portfolio/stocks.json?size=1000&pid=%d&category=%s&type=%s',
            'poster': 'https://xueqiu.com/service/poster',
            'group_list': 'https://xueqiu.com/v4/stock/portfolio/list.json?system=true'
        }
        self.group_infos = {}
        self.defalut_group_name = kwargs.get('group_name', '')
        self.defalut_group_info = {
             'category': 2,
             'type': 1,
             'id': -1,
             'name': ''
        }
        self.is_logon = False

    # 登录并获取分组信息
    def login(self):
        login_post_data = {
            'username': self.account_config.get('username', ''),
            'areacode': '86',
            'telephone': self.account_config['account'],
            'remember_me': '0',
            'password': self.account_config['password']
        }
        login_response = self.session.post(self.config['login_api'], data=login_post_data)
        login_status = json.loads(login_response.text)
        if 'error_description' in login_status:
            return False
        self.update_group_infos()
        self.is_logon = True
        return True

    # 获取一个分组名的信息dict
    def get_group_params(self, group_name):
        if group_name is None:
            group_name = self.defalut_group_name
        if isinstance(group_name,str):
            group_name = group_name.decode('utf-8')
        if group_name != '':
            if group_name not in self.group_infos.keys():
                self.create_group(group_name)  # 创建分组
                self.update_group_infos()  # 更新分组信息
                if group_name not in self.group_infos.keys():
                    return self.defalut_group_info
        return self.group_infos.get(group_name, self.defalut_group_info)

    # 获取一个分组下的自选股
    def get_stocks(self, group_name=None):
        d = self.get_group_params(group_name)
        stocks = self.session.get(self.config['stocks'] % (
            d.get('id', -1), d.get('category', 1), d.get('type', 1))).text
        stocks = json.loads(stocks)[u'stocks']
        return [x['code'] for x in stocks if (x['exchange'] == 'SH') or (x['exchange'] == 'SZ')]

    # 更新分组信息
    def update_group_infos(self):
        res = self.session.get(self.config['group_list']).text
        res = json.loads(res)[u'portfolios']
        self.group_infos = {}
        for p in res:
            self.group_infos[p['name']] = {'category': p['portfolio']['category'],
                                           'type': p['type'],
                                           'id': p['id'],
                                           'name': p['name']}

    # 创建一个分组
    def create_group(self, group_name):
        post_data = {
            'data[_]': get_unix_time(),
            'data[pname]': group_name,
            'url': '/stock/portfolio/create.json'
        }
        res = self.session.post(self.config['poster'], data=post_data).text
        return json.loads(res)['success']

    # 删除一个分组下的股票 group_name=None时，用参数里的默认分组。==''时为不分组添加
    def del_stock(self, stock, group_name=None):
        post_data = {
            'data[_]': get_unix_time(),
            'data[code]': stock,
            'url': '/stock/portfolio/delstock.json'
        }
        d = self.get_group_params(group_name)
        if d.get('name', '') != '':
            post_data['data[pname]'] = d.get('name', '')
            post_data['data[pids]'] = d.get('type', '')
        res = self.session.post(self.config['poster'], data=post_data).text
        return json.loads(res)['success']

    # 向一个分组下添加股票
    def add_stock(self, stock, group_name=None):
        post_data = {
            'data[_]': get_unix_time(),
            'data[code]': stock,
            'data[isnotice]': 1,
            'data[targetpercent]': 7,
            'url': '/stock/portfolio/addstock.json'
        }
        d = self.get_group_params(group_name)
        if d.get('name', '') != '':
            post_data['data[pnames]'] = d.get('name', '')
            post_data['data[pids]'] = d.get('type', '')
        res = self.session.post(self.config['poster'], data=post_data).text
        return json.loads(res)['success']

    def sync(self, stocks, group_name=None):
        """
        同步一个自选股列表到雪球自选股
        :param stocks: 雪球格式的自选股列表
        :param group_name: 雪球自选股分组名(不存在则自动创建)
        :return: 成功返回 '' 失败返回错误信息
        """
        try:
            if not self.is_logon:
                if not self.login():
                    return u'登录雪球错误'
            stocks = [chang_xq_stocks(x) for x in stocks] # 转成雪球的股票代码列表
            old_s = self.get_stocks(group_name) # 获取雪球分组自选股列表            
            for stock in [x for x in old_s if x not in stocks]:
                self.del_stock(stock, group_name)
            for stock in [x for x in stocks if x not in old_s]:
                self.add_stock(stock, group_name)
            return ''
        except Exception as e:
            return (u'同步雪球错误: %s' % str(e))

# 创建对象
xq = XqStocks(user_name='' # 用户名，一般为''
              , account='XXXXXXXXX' # 登录帐号，一般为手机号
              , password='XXXXXX' # 密码
              , group_name='测试' # 默认操作的自选股的分组名，默认为''
#               , proxy_ip='XXX.XXX.XXX.XXX' # http代码的IP # 不用代理可以不写
#               , proxy_port='808'  # http代理的端口
#               , proxy_user='' # 代理用户名,无需登录的用户名可以不写
#               , proxy_pwd='' # 代理密码
             )
# 单独测试
if not xq.login():
    print '登录错误'
else:
    stocks = xq.get_stocks()
    print '雪球自选股：'
    print stocks
    print '添加股票结果:%s' % xq.add_stock('SH000016')
    print '删除股票结果:%s' % xq.del_stock('SH600392')
    print '测试后雪球自选股:'
    print xq.get_stocks()

    # 同步测试
    stocks = ['000001.XSHE','399678.XSHE','600392.XSHG']
    xq.sync(stocks) # 将雪球自选股列表更新为stocks
    print '同步后雪球自选股:'
    print xq.get_stocks()

    print '获取所有雪球自选股列表'
    print xq.get_stocks('全部')