# -*- coding: utf-8 -*-
"""AIGU xtquant 兼容层 —— 现有 xtquant 代码零改造迁移

用法（把原来的 `from xtquant import xtdata` 换成下面两行即可）：

    import aigu_xt as xtdata
    xtdata.connect(api_key=os.environ["AIGU_API_KEY"])   # Key 从环境变量读取
    xtdata.get_full_tick(["600519.SH"])                   # 接口签名与 xtquant 相同

支持的接口：
    connect / get_stock_list_in_sector / get_full_tick
    get_market_data_ex(日K/分钟K) / subscribe_quote / unsubscribe_quote

与 xtquant 的差异：
    - connect 的 port 参数忽略；连接的是 AIGU 数据服务（Key 走环境变量或参数）
    - K线走 REST，冷门标的分钟线首次会触发服务端补数
    - subscribe_quote 回调参数为 {code: tick}（需 pip install websocket-client）
"""
import json
import os
import threading
import time
import urllib.parse
import urllib.request

# 数据服务地址（走 aiguq.com 域名反代，隐藏后端真实IP；白名单防SSRF）
_ALLOWED_HOSTS = {'aiguq.com', 'www.aiguq.com'}
BASE = 'http://aiguq.com'
_API_KEY = os.environ.get('AIGU_API_KEY', '')
_ws_subs = {}      # code -> (seq, callback)
_ws_seq = 0
_lock = threading.Lock()


def _check_url(url):
    p = urllib.parse.urlparse(url)
    if p.scheme not in ('http', 'https') or p.hostname not in _ALLOWED_HOSTS:
        raise ValueError('blocked url: ' + url)


def connect(port=None, ip=None, api_key=None):
    """兼容 xtdata.connect 签名；api_key 优先取参数，其次环境变量 AIGU_API_KEY"""
    global _API_KEY
    if api_key:
        _API_KEY = api_key
    if not _API_KEY:
        raise ValueError('缺少 API Key：connect(api_key=...) 或设环境变量 AIGU_API_KEY')
    url = BASE + '/healthz'
    _check_url(url)
    with urllib.request.urlopen(url, timeout=15) as resp:
        r = json.loads(resp.read().decode('utf-8', 'replace'))
    if not r.get('ok'):
        raise ConnectionError('AIGU 数据服务不可达: ' + BASE)
    return True


def _req(path, data=None, timeout=30):
    url = BASE + path + ('&key=' if '?' in path else '?key=') + _API_KEY
    _check_url(url)
    req = urllib.request.Request(url,
                                 data=json.dumps(data).encode() if data else None,
                                 headers={'X-API-Key': _API_KEY,
                                          'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = json.loads(resp.read().decode('utf-8', 'replace'))
    if body.get('code') != 0:
        raise RuntimeError(body.get('msg', 'API错误'))
    return body.get('data') or {}


def get_stock_list_in_sector(sector='沪深A股'):
    market = 'ALL' if '沪深' in sector or 'A股' in sector else sector
    d = _req('/v1/stock_list?market=%s' % market, timeout=60)
    return d.get('items') or []


def _norm(code):
    code = str(code)
    if '.' not in code and len(code) == 6 and code.isdigit():
        return code + ('.SH' if code.startswith(('6', '5', '9')) else '.SZ')
    return code


def get_full_tick(code_list):
    codes = ','.join(_norm(c) for c in code_list)
    d = _req('/v1/quote/batch?codes=' + codes, timeout=30)
    out = {}
    for it in d.get('items') or []:
        c = it['code']
        out[c] = {
            'time': it.get('time'), 'lastPrice': it.get('last'),
            'open': it.get('open'), 'high': it.get('high'), 'low': it.get('low'),
            'lastClose': it.get('pre_close'), 'volume': it.get('volume'),
            'amount': it.get('amount'), 'askPrice': it.get('ask'),
            'bidPrice': it.get('bid'), 'askVol': it.get('ask_vol'),
            'bidVol': it.get('bid_vol'), 'stockStatus': it.get('status'),
        }
    return out


def get_market_data_ex(field_list, stock_list, period='1d', start_time=None,
                       end_time=None, count=-1, dividend_type='none',
                       fill_data=True):
    """返回 {code: [bar字典...]}（轻量兼容，bar 含 ohlcv 字段）"""
    out = {}
    for c in stock_list:
        n = 120 if count in (-1, 0) else min(1000, count)
        d = _req('/v1/kline?code=%s&period=%s&count=%d' % (_norm(c), period, n),
                 timeout=60)
        bars = d.get('bars') or []
        out[str(c) if '.' in str(c) else _norm(c)] = [
            {'time': b['t'], 'open': b['o'], 'high': b['h'], 'low': b['l'],
             'close': b['c'], 'volume': b['v'], 'amount': b['a']}
            for b in bars]
    return out


# ── WebSocket 实时订阅 ───────────────────────────────
def subscribe_quote(stock_code, period='tick', callback=None, count=-1):
    """WebSocket 推送；回调收到 {code: tick}。需 pip install websocket-client"""
    if period != 'tick':
        return -1
    try:
        import websocket   # noqa: pip install websocket-client
    except ImportError:
        raise ImportError('实时订阅需安装: pip install websocket-client')

    global _ws_seq
    code = _norm(stock_code)
    url = BASE + '/ws?key=' + _API_KEY
    _check_url(url.replace('http', 'ws', 1) if url.startswith('http') else url)
    ws = websocket.WebSocketApp(
        url,
        on_open=lambda w: w.send(json.dumps({'op': 'sub', 'codes': [code]})),
        on_message=lambda w, m: _ws_dispatch(code, callback, m),
    )
    threading.Thread(target=ws.run_forever, daemon=True,
                     name='aigu-ws-' + code).start()
    with _lock:
        _ws_seq += 1
        seq = _ws_seq
        _ws_subs[seq] = (code, ws)
    return seq


def _ws_dispatch(code, callback, msg):
    try:
        d = json.loads(msg)
        if d.get('code') != code:
            return
        tick = {'time': d.get('time'), 'lastPrice': d.get('last'),
                'volume': d.get('volume')}
        callback({code: tick})
    except Exception:
        pass


def unsubscribe_quote(seq):
    with _lock:
        rec = _ws_subs.pop(int(seq), None)
    if rec:
        try:
            rec[1].close()
        except Exception:
            pass


def enable_hello(v=True):
    pass


def run():
    while True:
        time.sleep(1)