#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
build_long_history_seed.py
=============================================================================
长对话历史 90 天 / 12,000 条压测种子生成脚本
用于验证 P0 性能预算：
  1. 打开包含 12,000 条消息的长会话，P95 首屏加载至落底 ≤ 300ms
  2. 窗口只拉取最近 7 天 (可配置)，瞬时落底无平滑滚动动画过程
  3. 单日超额截断 (单日 500 条) 与 DayDivider 截断标记验证
  4. 触顶滚动补偿平滑加载更早日期分片
=============================================================================
"""

import sys
import os
import io
import json
import random
import subprocess
from datetime import datetime, timedelta, timezone

SESSION_ID = "sess_benchmark_long_001"
DAYS_COUNT = 90
TOTAL_TARGET = 12000

SAMPLE_INTENTS = [
    ("SUPPLEMENT_INFO", "consult", "user"),
    ("QUERY_PROGRESS", "consult", "user"),
    ("QUERY_REGULATIONS", "consult", "user"),
    ("CASUAL_CHAT", "casual", "user"),
    ("TRIGGER_TASK", "task", "user"),
    ("DOCGOV_REVIEW", "review", "user"),
]

SAMPLE_USER_PROMPTS = [
    "我们准备立项一款无线便携超声设备，请评估注册分类与型检周期。",
    "请查询 GB 9706.1-2020 关于患者漏电流和绝缘阻抗的强制性条款。",
    "超声探头手柄外壳采用什么等级的医用 PC 塑料能通过生物相容性？",
    "帮我对比一下国内迈瑞和飞利浦同类掌超产品的技术指标与主要卖点。",
    "关于电气安全和电磁兼容（EMC），型检摸底测试通常需要多长时间？",
    "项目进入可行性评审阶段，请同步准备初审材料和研发工时预算表。",
    "收到，下午两点我们与检验所专家开视频会对齐抽检样机数量。",
    "请帮我梳理一份向高管汇报的 5 分钟总结，重点突出风险和预算缺口。",
    "好的，没问题，先按既定计划推进第二轮软件单元测试。",
    "辛苦尽快起草一份注册路径初筛评估报告，我等下要向质量负责人汇报。",
]

SAMPLE_ASST_REPLIES = [
    "收到。便携无线超声诊断仪根据《医疗器械分类目录》，通常划分为第二类或第三类医疗器械（代码 07-01）。型检周期预估为 4-6 个月。",
    "根据 GB 9706.1-2020 标准第 8.7 条规定，应用部分在正常状态下对地漏电流限值为 100μA，单一故障状态下为 500μA。",
    "建议选用符合 ISO 10993-1 生物相容性标准的医用级聚碳酸酯（PC，如 Covestro Makrolon 2858），需完成细胞毒性、致敏与皮内刺激测试。",
    "已为您完成竞品多维比对：标杆竞品主机重量约 220g，单次续航 4 小时，支持 B/M/Color 多普勒模式；我们自研方案在功耗散热上具备明显优势。",
    "建议预留 3-4 周进行摸底整改，重点关注 YY 9706.102-2021 中静电放电（ESD）±8kV 接触放电与辐射发射（RE）余量。",
    "已为您在项目协作区生成《数字化便携超声注册路径可行性报告（初稿）》，关键排期节点已同步更新至项目协同看板。",
    "好的，会议要点已自动记录至协同黑板，稍后将为您生成会议纪要与待办任务清单。",
]


def generate_seed_csv():
    # 统一使用中国标准时区 (Asia/Shanghai, UTC+8)
    tz = timezone(timedelta(hours=8))
    now = datetime.now(tz)
    start_date = now - timedelta(days=DAYS_COUNT - 1)

    print(f"🚀 开始生成 90 天压测种子数据...")
    print(f"   会话 ID: {SESSION_ID}")
    print(f"   日期范围: {start_date.strftime('%Y-%m-%d')} ~ {now.strftime('%Y-%m-%d')}")
    print(f"   目标总数: 约 {TOTAL_TARGET} 条")

    # 确定每日消息条数：
    # 普通日：110~130 条
    # 高峰日（如 Day 30, Day 60, Day 85）：450~550 条（触发 >300 截断）
    day_counts = []
    for d in range(DAYS_COUNT):
        if d in (30, 60, 85):
            day_counts.append(random.randint(450, 520))
        else:
            day_counts.append(random.randint(110, 135))

    total_planned = sum(day_counts)
    print(f"   实际规划条数: {total_planned} 条 (包含 3 个单日超 500 条的高频爆量日)")

    output = io.StringIO()
    msg_id = 1

    for d, count in enumerate(day_counts):
        day_dt = start_date + timedelta(days=d)
        day_key = day_dt.strftime("%Y-%m-%d")

        for i in range(count):
            # 时间按当天工作时间均匀分布 (08:30 ~ 21:30)
            seconds_offset = int((i / max(1, count)) * 46800)
            msg_dt = day_dt.replace(hour=8, minute=30, second=0) + timedelta(seconds=seconds_offset)
            ts_str = msg_dt.strftime("%Y-%m-%d %H:%M:%S%z")

            # 偶数用户，奇数助手
            if i % 2 == 0:
                role = "user"
                uid_str = "1"
                intent, category, _ = random.choice(SAMPLE_INTENTS)
                content = random.choice(SAMPLE_USER_PROMPTS) + f" (d{d}-m{i})"
                meta = {"author_name": "张工"}
                if category == "file":
                    meta["attachments"] = [{"name": "需求规格说明书.pdf", "size": "3.5MB", "type": "pdf"}]
                elif d % 15 == 0 and i == 0:
                    category = "milestone"
                    meta["milestone"] = True
            else:
                role = "assistant"
                uid_str = r"\N"
                intent = "ADVISOR_BRIEF"
                category = "consult"
                content = random.choice(SAMPLE_ASST_REPLIES) + f" [参考编号 CDMO-{day_key}-{i:03d}]"
                meta = {"author_name": "项目助手"}

            meta_json = json.dumps(meta, ensure_ascii=False)

            # CSV: session_id, user_id, role, content, intent, branch, metadata, day_key, category, created_at
            # 使用制表符作为分隔符
            content_clean = content.replace("\t", " ").replace("\n", " ")
            meta_clean = meta_json.replace("\t", " ").replace("\n", " ")

            output.write(
                f"{SESSION_ID}\t{uid_str}\t{role}\t{content_clean}\t{intent}\t\t{meta_clean}\t{day_key}\t{category}\t{ts_str}\n"
            )
            msg_id += 1

    return output.getvalue(), total_planned


def import_to_postgres(csv_data, total_count):
    print(f"📦 正在通过 psql COPY 流式注入容器 shuimu-context-pg-dev...")

    # 1. 确保 chat_sessions 存在
    ensure_session_sql = f"""
    INSERT INTO chat_sessions (session_id, user_id, project_id, title, created_at, updated_at)
    VALUES ('{SESSION_ID}', 101, 1, '90天万条长历史压测基准会话', now(), now())
    ON CONFLICT (session_id) DO UPDATE SET updated_at = now();
    DELETE FROM chat_messages WHERE session_id = '{SESSION_ID}';
    """

    subprocess.run(
        [
            "docker", "exec", "-i", "shuimu-context-pg-dev",
            "psql", "-U", "app", "-d", "agent_system_dev", "-c", ensure_session_sql
        ],
        check=True,
        capture_output=True,
    )

    # 2. 执行 COPY 导入
    copy_sql = r"""
    COPY chat_messages (session_id, user_id, role, content, intent, branch, metadata, day_key, category, created_at) FROM STDIN WITH (FORMAT text, DELIMITER E'\t');
    """

    p = subprocess.Popen(
        [
            "docker", "exec", "-i", "shuimu-context-pg-dev",
            "psql", "-U", "app", "-d", "agent_system_dev", "-c", copy_sql
        ],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    stdout, stderr = p.communicate(input=csv_data)
    if p.returncode != 0:
        print(f"❌ 导入失败: {stderr}", file=sys.stderr)
        sys.exit(1)

    print(f"✅ 成功导入 {total_count} 条消息！")

    # 3. 统计验证
    verify_cmd = [
        "docker", "exec", "-i", "shuimu-context-pg-dev",
        "psql", "-U", "app", "-d", "agent_system_dev", "-c",
        f"""
        SELECT 
            count(*) as total_msgs,
            count(DISTINCT day_key) as total_days,
            min(day_key) as earliest_day,
            max(day_key) as latest_day
        FROM chat_messages 
        WHERE session_id = '{SESSION_ID}';
        """
    ]
    res = subprocess.check_output(verify_cmd, text=True)
    print("\n📊 数据库验证统计：")
    print(res)


if __name__ == "__main__":
    csv_data, total_planned = generate_seed_csv()
    import_to_postgres(csv_data, total_planned)
    print("🎉 90天万条长历史压测数据准备完毕！")