# -*- coding: utf-8 -*-
"""
AgentLoop Dataset 数据写入示例

演示通过 ExecuteQuery 接口执行 SQL 进行数据变更：
- INSERT   插入数据（单条 & 批量）
- UPDATE   按 id 主键更新
- DELETE   按 id 主键删除

所有写入操作统一使用 ExecuteQuery(type="SQL") 接口。

使用前请设置环境变量：
- SLSDEMO_ALIYUN_ACCESS_KEY_ID
- SLSDEMO_ALIYUN_ACCESS_KEY_SECRET
- ALIBABA_CLOUD_CMS_ENDPOINT   (如: cms.cn-shanghai.aliyuncs.com)
- ALIBABA_CLOUD_CMS_WORKSPACE
"""

import os

from alibabacloud_cms20240330.client import Client
from alibabacloud_cms20240330.models import ExecuteQueryRequest
from alibabacloud_tea_openapi.models import Config
from dotenv import load_dotenv

load_dotenv()

DATASET_NAME = "dataset_demo"


def get_client() -> Client:
    config = Config(
        access_key_id=os.getenv("SLSDEMO_ALIYUN_ACCESS_KEY_ID"),
        access_key_secret=os.getenv("SLSDEMO_ALIYUN_ACCESS_KEY_SECRET"),
        endpoint=os.getenv("ALIBABA_CLOUD_CMS_ENDPOINT"),
    )
    return Client(config)


def execute_sql(client: Client, workspace: str, sql: str) -> dict:
    """执行 SQL（INSERT / UPDATE / DELETE）"""
    request = ExecuteQueryRequest(query=sql, type="SQL")
    response = client.execute_query(workspace, DATASET_NAME, request)
    return response.body.to_map()


def print_result(operation: str, result: dict):
    meta = result["meta"]
    print(f"✓ {operation}, affected_rows: {meta['affectedRows']}, elapsed: {meta['elapsedMillisecond']}ms")


# =============================================================================
# INSERT: 插入单条数据
# =============================================================================

def example_insert_single(client: Client, workspace: str):
    """插入单条数据。"""
    sql = f"""
    INSERT INTO {DATASET_NAME} (input, output, model, score, metadata)
    VALUES (
        '如何查看最近一小时的错误日志？',
        '使用查询语句: level:ERROR | SELECT * FROM log WHERE __time__ > now() - interval ''1 hour''',
        'qwen-plus',
        0.95,
        '{{"input_tokens": 15, "output_tokens": 42}}'
    )
    """
    result = execute_sql(client, workspace, sql)
    print_result("单条插入", result)


# =============================================================================
# INSERT: 批量插入多条数据
# =============================================================================

def example_insert_batch(client: Client, workspace: str):
    """
    批量插入多条数据。

    一次 INSERT 可包含多组 VALUES，减少网络往返。
    """
    sql = f"""
    INSERT INTO {DATASET_NAME} (input, output, model, score, metadata)
    VALUES
        (
            '统计今天各接口的调用次数',
            'SELECT api_path, count(*) as cnt FROM log GROUP BY api_path ORDER BY cnt DESC',
            'gpt-4o',
            0.88,
            '{{"input_tokens": 12, "output_tokens": 35}}'
        ),
        (
            '查找响应时间超过5秒的请求',
            'latency > 5000 | SELECT request_id, latency, api_path FROM log',
            'claude-3.5-sonnet',
            0.92,
            '{{"input_tokens": 10, "output_tokens": 28}}'
        ),
        (
            '分析用户登录失败的原因分布',
            'action:login AND status:failed | SELECT error_reason, count(*) as cnt FROM log GROUP BY error_reason',
            'qwen-plus',
            0.85,
            '{{"input_tokens": 14, "output_tokens": 45}}'
        )
    """
    result = execute_sql(client, workspace, sql)
    print_result("批量插入 (3条)", result)


# =============================================================================
# UPDATE: 按 ID 更新数据
# =============================================================================

def example_update_by_id(client: Client, workspace: str, doc_id: str):
    """
    按 ID 更新单条数据。

    注意：UPDATE 必须通过 id 主键进行，不支持按条件批量更新。
    可先通过查询获取目标数据的 id，再执行更新。
    """
    sql = f"""
    UPDATE {DATASET_NAME}
    SET score = 0.98,
        output = '优化后的回答内容...'
    WHERE id = '{doc_id}'
    """
    result = execute_sql(client, workspace, sql)
    print_result(f"更新 id={doc_id}", result)


# =============================================================================
# DELETE: 按 ID 删除数据
# =============================================================================

def example_delete_by_id(client: Client, workspace: str, doc_id: str):
    """
    按 ID 删除单条数据。

    注意：DELETE 必须通过 id 主键进行，不支持按条件批量删除。
    """
    sql = f"DELETE FROM {DATASET_NAME} WHERE id = '{doc_id}'"
    result = execute_sql(client, workspace, sql)
    print_result(f"删除 id={doc_id}", result)


# =============================================================================
# 辅助: 查询获取 id（用于 UPDATE / DELETE 演示）
# =============================================================================

def get_first_doc_id(client: Client, workspace: str) -> str | None:
    """查询第一条数据的 id，用于后续 UPDATE / DELETE 演示。"""
    sql = f"SELECT id, input FROM {DATASET_NAME} LIMIT 1"
    request = ExecuteQueryRequest(query=sql, type="SQL")
    response = client.execute_query(workspace, DATASET_NAME, request)
    data = response.body.to_map().get("data", [])
    if data:
        print(f"  查询到 id={data[0]['id']}, input: {data[0].get('input', '')[:30]}...")
        return data[0]["id"]
    print("  无数据")
    return None


# =============================================================================
# 主函数
# =============================================================================

def main():
    client = get_client()
    workspace = os.getenv("ALIBABA_CLOUD_CMS_WORKSPACE")

    print(f"AgentLoop Dataset 数据写入示例 (dataset: {DATASET_NAME})")
    print("=" * 50)

    # INSERT
    print("\n【INSERT 单条】")
    example_insert_single(client, workspace)

    print("\n【INSERT 批量】")
    example_insert_batch(client, workspace)

    # UPDATE（先查 id 再更新）
    print("\n【UPDATE 按 ID】")
    doc_id = get_first_doc_id(client, workspace)
    if doc_id:
        example_update_by_id(client, workspace, doc_id)

    # DELETE（先查 id 再删除）
    print("\n【DELETE 按 ID】")
    doc_id = get_first_doc_id(client, workspace)
    if doc_id:
        example_delete_by_id(client, workspace, doc_id)

    print("\n" + "=" * 50)
    print("数据写入示例完成")


if __name__ == "__main__":
    main()