# -*- coding: utf-8 -*-
"""AI 助手「大结果集图表精炼 + Excel 导出」后端确定性构建。

产品口径（用户已拍板，勿自作主张）：
1. 先有汇总表（答案里的 markdown 表），再据表选图 —— 图表数据 100% 来自该表，不调 LLM、不编数字。
2. 仅在用户显式要求图表时才出图（关键词见 CHART_INTENT_KEYWORDS）；不带这类词的提问行为完全不变。
3. 报告形态 = Excel 导出（真 .xlsx，后端端点生成）。
4. 图表形态确定性选择：
     表内有时序列（年/季/月/周/日）→ 折线 line
     只有一个可加总数值列        → 环形 donut
     其余                        → 柱状 bar
   无可用数值列 → type='table'（只给表 + 导出按钮）。

该模块不依赖 ai_chat，可独立单测。
"""
import io
import os
import re
import datetime
import zipfile

# ---------------------------------------------------------------- 图表意图关键词
# 仅当问题里出现下列任一子串时，才在后处理阶段附加 chart/table 结构。
# 注意：刻意不收裸「趋势」（"行业趋势"是常规提问），只收「看趋势/趋势图」这类明确要图的说法。
CHART_INTENT_KEYWORDS = [
    '用图表', '用个图表', '上图表', '图表', '画图', '画个图', '画一张图', '画张图',
    '可视化', '图形展示', '图形化', '图形呈现', '可视化展示', '可视化呈现', '图形呈现',
    '给我看趋势', '看趋势', '看下趋势', '看看趋势', '趋势图', '折线图', '折线',
    '柱状图', '柱形图', '条形图', '环形图', '圆环图', '甜甜圈图', '饼图', '扇形图',
    '出个图', '来个图', '上个图', '展示成图', '用图显示', '用图展示', '出图',
]


def has_chart_intent(question):
    q = (question or '').strip()
    if not q:
        return False
    return any(k in q for k in CHART_INTENT_KEYWORDS)


# ---------------------------------------------------------------- markdown 表解析
def _is_md_row(line):
    s = (line or '').strip()
    return s.startswith('|') and s.count('|') >= 2


def _is_sep_row(line):
    s = (line or '').strip().strip('|')
    cells = [c.strip() for c in s.split('|')]
    if not cells:
        return False
    return all(re.fullmatch(r':?-{1,}:?', c) for c in cells if c != '')


def _split_row(line):
    s = (line or '').strip()
    if s.startswith('|'):
        s = s[1:]
    if s.endswith('|'):
        s = s[:-1]
    return [c.strip() for c in s.split('|')]


def parse_md_tables(text):
    """从答案正文里抽出所有 markdown 表，返回 [{'headers':[...], 'rows':[[...]]}]。

    表格定义：连续 >=2 行且每行以 '|' 开头，第二行是分隔行（---）。"""
    tables = []
    lines = (text or '').split('\n')
    n = len(lines)
    i = 0
    while i < n:
        if _is_md_row(lines[i]):
            j = i
            block = []
            while j < n and _is_md_row(lines[j]):
                block.append(lines[j].strip())
                j += 1
            if len(block) >= 2 and _is_sep_row(block[1]):
                headers = _split_row(block[0])
                rows = [_split_row(r) for r in block[2:]]
                rows = [r for r in rows if any((c or '').strip() != '' for c in r)]
                # 补齐/裁剪到表头列数
                norm = []
                for r in rows:
                    if len(r) < len(headers):
                        r = r + [''] * (len(headers) - len(r))
                    elif len(r) > len(headers):
                        r = r[:len(headers)]
                    norm.append(r)
                if headers and norm:
                    tables.append({'headers': headers, 'rows': norm})
            i = j
        else:
            i += 1
    return tables


# ---------------------------------------------------------------- 数值解析
_NUM_UNITS = ('元', '个', '次', '条', '笔', '家', '台', '件', '人', '天', '小时',
              'h', 'kg', 'g', 'ml', 'L', '套', '份', '单', '万', '亿')


def _to_number(cell):
    """把单元格文本解析为数字（支持千分位、货币符、%、万/亿），失败返回 None。"""
    if cell is None:
        return None
    s = str(cell).strip()
    if s == '' or s in ('-', '—', '--', '/', 'N/A', 'NA', 'null', 'None', '无'):
        return None
    s = s.replace(',', '').replace('，', '')
    for ch in ('¥', '￥', '$', ' ', '\u3000'):
        s = s.replace(ch, '')
    mult = 1.0
    if s.endswith('亿'):
        mult = 1e8
        s = s[:-1]
    elif s.endswith('万'):
        mult = 1e4
        s = s[:-1]
    if s.endswith('%'):
        s = s[:-1]
    # 去尾部单位（在万/亿之后再剥一次，兼容 "12万元"）
    for _ in range(2):
        for u in _NUM_UNITS:
            if s.endswith(u) and u not in ('万', '亿'):
                s = s[:-len(u)]
                break
        else:
            break
    if s == '' or s == '-':
        return None
    try:
        v = float(s)
    except Exception:
        return None
    return v * mult


def _col_is_numeric(rows, c):
    vals = [(r[c] if c < len(r) else '') for r in rows]
    non_empty = [v for v in vals if str(v).strip() != '']
    if not non_empty:
        return False
    return all(_to_number(v) is not None for v in non_empty)


# ---------------------------------------------------------------- 时序列识别
_YEAR = re.compile(r'^\s*\d{4}\s*年?\s*$')
_MONTH = re.compile(r'^\s*\d{1,2}\s*月\s*$|^\s*\d{4}\s*[-/年]\s*\d{1,2}\s*月?\s*$|^\s*\d{4}\s*[-/]\s*\d{1,2}\s*$')
_DAY = re.compile(r'^\s*\d{4}\s*[-/]\s*\d{1,2}\s*[-/]\s*\d{1,2}\s*$|^\s*\d{1,2}\s*月\s*\d{1,2}\s*日?\s*$')
_WEEK = re.compile(r'第\s*\d{1,2}\s*周|^\s*\d{4}\s*[-/]?\s*[Ww]\s*\d{1,2}\s*$|^\s*[Ww]\s*\d{1,2}\s*$|^\s*\d{1,2}\s*周\s*$')
_QUARTER = re.compile(r'^\s*(\d{4}\s*年?\s*)?[Qq]?[1-4]\s*季度\s*$|^\s*(\d{4}\s*年?\s*)?[Qq][1-4]\s*$')


def _time_series_kind(labels):
    labs = [str(x).strip() for x in labels if str(x).strip() != '']
    if len(labs) < 2:
        return None
    for kind, rx in (('day', _DAY), ('week', _WEEK), ('month', _MONTH),
                     ('quarter', _QUARTER), ('year', _YEAR)):
        hit = sum(1 for x in labs if rx.search(x))
        if hit >= max(2, int(len(labs) * 0.6 + 0.999)):
            return kind
    return None


# ---------------------------------------------------------------- 权限过滤
def _core(name):
    """客户核心词归一（与 ai_chat._cust_core 口径一致；此处内联避免循环依赖）。"""
    s = (name or '').strip()
    if not s:
        return ''
    for lb, rb in (('（', '）'), ('(', ')')):
        while lb in s and rb in s:
            i = s.find(lb)
            j = s.find(rb, i)
            if j > i:
                s = s[:i] + s[j + 1:]
            else:
                break
    for suf in ('股份有限公司', '有限责任公司', '有限公司', '公司', '集团'):
        if s.endswith(suf) and len(s) - len(suf) >= 2:
            s = s[:-len(suf)]
    changed = True
    while changed:
        changed = False
        for suf in ('科技', '医疗', '医药', '医学', '技术', '制药', '工程', '诊断', '检测', '生物'):
            if s.endswith(suf) and len(s) - len(suf) >= 2:
                if suf == '生物' and len(s) - len(suf) < 3:
                    continue
                s = s[:-len(suf)]
                changed = True
    return s


def _perm_filter_rows(rows, all_cores, visible_cores):
    """只保留该用户可见的行：首列命中客户名且核心词不在可见集 → 丢弃。"""
    if not all_cores or visible_cores is None:
        return rows, 0
    kept, dropped = [], 0
    for r in rows:
        first = _core(r[0] if r else '')
        if first and first in all_cores and first not in visible_cores:
            dropped += 1
            continue
        kept.append(r)
    return kept, dropped


# ---------------------------------------------------------------- 组图
def _series_cols(headers, rows):
    return [c for c in range(1, len(headers)) if _col_is_numeric(rows, c)]


def _find_title(answer, table):
    """取表格上方最近的 markdown 标题/加粗行作为图题。"""
    lines = (answer or '').split('\n')
    idx = None
    for i, ln in enumerate(lines):
        if _is_md_row(ln):
            # 确认是我们要的那张表（表头匹配）
            hdr = _split_row(ln)
            if hdr[:1] == table['headers'][:1]:
                idx = i
                break
            idx = i
    if idx is None:
        return '数据汇总'
    for k in range(idx - 1, max(-1, idx - 6), -1):
        t = (lines[k] or '').strip()
        if not t:
            continue
        t = re.sub(r'^#{1,6}\s*', '', t)
        t = re.sub(r'^\*\*(.*)\*\*$', r'\1', t)
        t = t.strip('*#：: ')
        if t:
            return t[:40]
    return '数据汇总'


def build_chart(question, answer, all_cores=None, visible_cores=None):
    """产出确定性图表结构；无意图/无可解析表 → None。"""
    if not has_chart_intent(question):
        return None
    tables = parse_md_tables(answer)
    if not tables:
        return None

    # 选表：优先第一张「可成图」的表（>=1 数值列），否则第一张 >=2 行的表
    chosen = None
    for t in tables:
        if len(t['rows']) >= 2 and len(t['headers']) >= 2:
            if _series_cols(t['headers'], t['rows']):
                chosen = t
                break
            if chosen is None:
                chosen = t
    if chosen is None:
        chosen = tables[0]

    headers = chosen['headers']
    rows, dropped = _perm_filter_rows(chosen['rows'], all_cores, visible_cores)

    labels = [(r[0] if r else '') for r in rows]
    sercols = _series_cols(headers, rows)
    kind = _time_series_kind(labels)

    if sercols and kind:
        ctype = 'line'
    elif len(sercols) == 1:
        ctype = 'donut'
    elif sercols:
        ctype = 'bar'
    else:
        ctype = 'table'

    series = []
    for c in sercols:
        series.append({
            'name': headers[c],
            'values': [_to_number(r[c] if c < len(r) else '') for r in rows],
        })

    payload = {
        'type': ctype,
        'title': _find_title(answer, chosen),
        'labels': [str(x) for x in labels],
        'series': series,
        'table': {'headers': list(headers), 'rows': [list(r) for r in rows]},
        'x_title': headers[0] if headers else '',
        'y_title': series[0]['name'] if series else '',
        'dropped_rows': dropped,
        'time_kind': kind or '',
        'source': 'answer_markdown_table',
    }
    return payload


# ---------------------------------------------------------------- Excel 导出
def _safe_filename(name):
    s = re.sub(r'[\\/:*?"<>|\r\n\t]+', '_', (name or '数据汇总')).strip() or '数据汇总'
    return s[:60]


def build_xlsx(title, headers, rows):
    """生成真 .xlsx 字节流。首行 == 表头，行数 == len(rows)+1（供验收核对）。"""
    headers = [str(h) for h in (headers or [])]
    rows = rows or []
    sheet = _safe_filename(title or '数据汇总')[:31] or '数据汇总'
    try:
        import openpyxl
        from openpyxl.styles import Font, PatternFill, Alignment
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = sheet
        ws.append(headers)
        for c in range(1, len(headers) + 1):
            cell = ws.cell(row=1, column=c)
            cell.font = Font(bold=True, color='FFFFFF')
            cell.fill = PatternFill('solid', fgColor='1A56DB')
            cell.alignment = Alignment(horizontal='center', vertical='center')
        for r in rows:
            out = []
            for i, v in enumerate(r):
                num = _to_number(v)
                out.append(num if num is not None else str(v))
            ws.append(out)
        for c in range(1, len(headers) + 1):
            ws.column_dimensions[openpyxl.utils.get_column_letter(c)].width = 16
        buf = io.BytesIO()
        wb.save(buf)
        return buf.getvalue()
    except Exception:
        return _build_xlsx_minimal(headers, rows, sheet)


def _xml_esc(s):
    return (str(s).replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
            .replace('"', '&quot;'))


def _col_letter(n):
    s = ''
    while n > 0:
        n, r = divmod(n - 1, 26)
        s = chr(65 + r) + s
    return s


def _build_xlsx_minimal(headers, rows, sheet):
    """无 openpyxl 时的兜底：手写最小 OOXML（仅 inlineStr/数值）。"""
    def cell_xml(ref, val):
        num = _to_number(val)
        if num is not None:
            return '<c r="%s"><v>%s</v></c>' % (ref, repr(num) if isinstance(num, float) else num)
        return '<c r="%s" t="inlineStr"><is><t xml:space="preserve">%s</t></is></c>' % (ref, _xml_esc(val))

    def row_xml(ridx, vals):
        cells = ''.join(cell_xml('%s%d' % (_col_letter(i + 1), ridx), v) for i, v in enumerate(vals))
        return '<row r="%d">%s</row>' % (ridx, cells)

    sheet_rows = [row_xml(1, headers)]
    for i, r in enumerate(rows):
        sheet_rows.append(row_xml(i + 2, r))
    sheet_xml = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
                 '<worksheet xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main">'
                 '<sheetData>%s</sheetData></worksheet>' % ''.join(sheet_rows))
    ct = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
          '<Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types">'
          '<Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/>'
          '<Default Extension="xml" ContentType="application/xml"/>'
          '<Override PartName="/xl/workbook.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml"/>'
          '<Override PartName="/xl/worksheets/sheet1.xml" ContentType="application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml"/>'
          '</Types>')
    rels = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
            '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
            '<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="xl/workbook.xml"/>'
            '</Relationships>')
    wb = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
          '<workbook xmlns="http://schemas.openxmlformats.org/spreadsheetml/2006/main" '
          'xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships">'
          '<sheets><sheet name="%s" sheetId="1" r:id="rId1"/></sheets></workbook>' % _xml_esc(sheet))
    wbrels = ('<?xml version="1.0" encoding="UTF-8" standalone="yes"?>'
              '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">'
              '<Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet" Target="worksheets/sheet1.xml"/>'
              '</Relationships>')
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as z:
        z.writestr('[Content_Types].xml', ct)
        z.writestr('_rels/.rels', rels)
        z.writestr('xl/workbook.xml', wb)
        z.writestr('xl/_rels/workbook.xml.rels', wbrels)
        z.writestr('xl/worksheets/sheet1.xml', sheet_xml)
    return buf.getvalue()


def export_filename(title):
    d = datetime.datetime.now().strftime('%Y%m%d')
    return '%s_%s.xlsx' % (_safe_filename(title), d)