"""
独立测试脚本 — 使用外部 LLM API 驱动一个完整的单日交易 episode。

直接调用 OpenAI-compatible API 驱动模型完成一次交易日

用途：
- 验证整个 Agent 循环是否跑通（system prompt → 数据工具 → TakeAction）
- 观察模型输出的消息格式和质量
- 快速调试 prompt / 环境 / 工具链

Usage (需在训练环境 verl-next conda 中运行):
    python scripts/test_agent.py --base-url http://localhost:8000/v1 --model Qwen/Qwen3-4B --date 2025-03-15 --universe SH600519 SH601398 SH600036
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Any

import openai
from openai.types.chat import ChatCompletionToolParam
from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall

# 确保项目根目录在 sys.path 中
sys.path.insert(0, str(Path(__file__).parent.parent))

from trading_tools import (
    EnvConfig,
    Portfolio,
    compute_alpha_factor,
    ensure_qlib,
    format_observation,
    get_current_price,
    get_macro_indicators,
    get_prev_close,
    get_price_data,
    update_portfolio_prices,
)
from market_state import MarketStateCalculator


# 配置 logging：只输出到 stderr，不干扰 stdout 的消息展示
logger = logging.getLogger(__name__)
handler = logging.StreamHandler(sys.stderr)
handler.setLevel(logging.INFO)
formatter = logging.Formatter("%(asctime)s | %(levelname)-7s | %(message)s", datefmt="%H:%M:%S")
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)


def build_tools_schema() -> list[ChatCompletionToolParam]:
    """构建 OpenAI function calling 的工具定义。"""
    return [
        {
            "type": "function",
            "function": {
                "name": "get_price_data",
                "description": "获取单个标的的历史 K 线数据和技术指标",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "symbol": {"type": "string", "description": "股票代码，如 SH600519"},
                        "period": {"type": "string", "description": "时间跨度，如 20d、60d"},
                        "indicators": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "技术指标列表，如 [\"MA5\", \"RSI_14\"]",
                        },
                    },
                    "required": ["symbol", "period", "indicators"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "get_macro_indicators",
                "description": "获取大盘指数行情数据",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "index": {"type": "string", "description": "指数代码，默认 SH000300"},
                        "period": {"type": "string", "description": "时间跨度，如 20d"},
                    },
                    "required": [],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "compute_alpha_factor",
                "description": "使用 Qlib 表达式引擎计算自定义 Alpha 因子",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "expression": {"type": "string", "description": "Qlib 表达式，如 $close / Ref($close, 5) - 1"},
                        "symbols": {
                            "type": "array",
                            "items": {"type": "string"},
                            "description": "标的列表",
                        },
                        "period": {"type": "string", "description": "时间跨度，默认 60d"},
                    },
                    "required": ["expression", "symbols"],
                },
            },
        },
        {
            "type": "function",
            "function": {
                "name": "TakeAction",
                "description": "Agent 的交易决策（buy / sell / stop）",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "type": {"type": "string", "enum": ["buy", "sell", "stop"]},
                        "symbol": {"type": "string", "description": "股票代码"},
                        "amount": {"type": "integer", "description": "交易股数（100的整数倍）"},
                        "reasoning": {"type": "string", "description": "决策理由"},
                    },
                    "required": ["type"],
                },
            },
        },
    ]


def run_episode(
    *,
    base_url: str,
    api_key: str,
    model: str,
    date: str,
    universe: list[str],
    initial_cash: float,
    qlib_data_dir: str,
    max_turns: int,
    enable_thinking: bool,
) -> None:
    """运行一个完整的单日交易 episode。"""
    # --- 初始化 Qlib + Portfolio ---
    ensure_qlib(qlib_data_dir)
    env_config = EnvConfig(qlib_data_dir=qlib_data_dir)

    portfolio = Portfolio(
        initial_cash=initial_cash,
        open_cost=env_config.open_cost,
        close_cost=env_config.close_cost,
        min_cost=env_config.min_cost,
        slippage_rate=env_config.slippage_rate,
    )
    actions_this_day = 0
    max_actions_per_day = len(universe) if universe else env_config.max_actions_per_day
    done = False

    # --- 构建 system prompt ---
    market_calc = MarketStateCalculator(env_config)
    try:
        macro_insights = market_calc.calculate(date).to_text()
    except Exception as e:
        logger.warning(f"macro_insights 计算失败: {e}")
        macro_insights = ""

    system_prompt = (
        f"你是一个 A 股交易助手。今日日期：{date}。\n"
        f"股票池：{', '.join(universe)}\n"
        f"{macro_insights}"
    )

    initial_obs_text = format_observation(
        portfolio=portfolio,
        current_date=date,
        universe=universe,
        actions_this_day=0,
        max_actions_per_day=max_actions_per_day,
        done=False,
    )
    user_msg = (
        f"当前交易日：{date}。\n"
        f"初始持仓：空仓\n"
        f"可用现金：{initial_cash:,.2f} 元\n\n"
        f"请开始分析市场并做出交易决策。\n\n"
        f"---\n\n{initial_obs_text}"
    )

    # --- 初始化 LLM client ---
    client = openai.OpenAI(base_url=base_url, api_key=api_key)
    tools_schema = build_tools_schema()

    messages: list[Any] = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_msg},
    ]

    print("\n" + "=" * 80)
    print(f" 交易日: {date} | 股票池: {', '.join(universe)} | 初始资金: {initial_cash:,.0f}")
    print("=" * 80)
    print(f"\n[SYSTEM PROMPT] ({len(system_prompt)} chars)")
    print(f"\n[USER MESSAGE]\n{user_msg}")

    # --- Agent loop ---
    turn = 0

    while not done and turn < max_turns:
        turn += 1
        print(f"\n{'─' * 60}")
        print(f" Turn {turn}/{max_turns}")
        print(f"{'─' * 60}")

        extra_body = (
            {"chat_template_kwargs": {"enable_thinking": True}}
            if enable_thinking
            else None
        )

        try:
            response = client.chat.completions.create(
                model=model,
                messages=messages,
                tools=tools_schema,
                tool_choice="auto",
                extra_body=extra_body,
            )
        except Exception as e:
            logger.error(f"LLM API 调用失败: {e}")
            break

        choice = response.choices[0]
        assistant_msg = choice.message

        if assistant_msg.content:
            print(f"\n[ASSISTANT]\n{assistant_msg.content}")

        if not assistant_msg.tool_calls:
            print("\n[INFO] 模型未调用工具，episode 结束")
            break

        func_tool_calls = [
            tc for tc in assistant_msg.tool_calls
            if isinstance(tc, ChatCompletionMessageToolCall)
        ]
        msg_dict: dict[str, Any] = {"role": "assistant", "content": assistant_msg.content or ""}
        msg_dict["tool_calls"] = [
            {
                "id": tc.id,
                "type": "function",
                "function": {"name": tc.function.name, "arguments": tc.function.arguments},
            }
            for tc in func_tool_calls
        ]
        messages.append(msg_dict)

        for tc in func_tool_calls:
            func_name = tc.function.name
            try:
                args = json.loads(tc.function.arguments)
            except json.JSONDecodeError:
                args = {}

            print(f"\n[TOOL CALL] {func_name}({json.dumps(args, ensure_ascii=False)})")

            if func_name == "TakeAction":
                action_type = args["type"]
                symbol = args.get("symbol")
                amount = args.get("amount")
                reasoning = args.get("reasoning", "")

                if action_type in ("buy", "sell") and symbol and amount:
                    if universe and symbol not in universe:
                        logger.warning(f"{symbol} 不在股票池中，交易被拒绝")
                    else:
                        current_price = get_current_price(symbol, date)
                        if current_price is None:
                            logger.warning(f"无法获取 {symbol} 当前价格")
                        else:
                            prev_close = get_prev_close(symbol, date)
                            record = portfolio.execute_trade(
                                symbol=symbol,
                                direction=action_type,
                                shares=amount,
                                current_price=current_price,
                                prev_close=prev_close or current_price,
                                current_date=date,
                            )
                            if record is not None:
                                record.reasoning = reasoning
                                logger.info(f"交易执行: {action_type} {symbol} {amount}股 @ {record.price:.2f}")
                            else:
                                logger.info(f"交易被拒绝: {action_type} {symbol} {amount}股")

                    actions_this_day += 1
                    if actions_this_day >= max_actions_per_day:
                        done = True

                elif action_type == "stop":
                    logger.info(f"stop 决策: {reasoning}")
                    done = True

                # 生成观测文本
                if done:
                    update_portfolio_prices(portfolio, date)

                obs_text = format_observation(
                    portfolio=portfolio,
                    current_date=date,
                    universe=universe,
                    actions_this_day=actions_this_day,
                    max_actions_per_day=max_actions_per_day,
                    done=done,
                )

                if action_type == "stop":
                    result_text = f"[交易日结束] {reasoning}\n\n{obs_text}"
                elif action_type in ("buy", "sell"):
                    result_text = (
                        f"[{action_type.upper()}] {symbol} {amount}股\n"
                        f"理由: {reasoning}\n\n{obs_text}"
                    )
                else:
                    result_text = obs_text

                print(f"\n[TOOL RESULT] (done={done})\n{result_text[:500]}")

                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result_text,
                })

                if done:
                    break
            else:
                # 数据工具 — 直接调用对应函数
                try:
                    if func_name == "get_price_data":
                        result_text = get_price_data(date, **args)
                    elif func_name == "get_macro_indicators":
                        result_text = get_macro_indicators(date, **args)
                    elif func_name == "compute_alpha_factor":
                        result_text = compute_alpha_factor(date, **args)
                    else:
                        result_text = f"[错误] 未知工具: {func_name}"
                except Exception as e:
                    result_text = f"[错误] {e}"

                print(f"\n[TOOL RESULT]\n{result_text[:500]}{'...' if len(result_text) > 500 else ''}")

                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result_text,
                })

    # --- Episode 结束 ---
    print(f"\n{'=' * 80}")
    print(f" Episode 完成 | Turns: {turn} | Done: {done}")
    print(f"{'=' * 80}")

    nav = portfolio.get_nav()
    cumulative_return = nav / portfolio.initial_cash - 1 if portfolio.initial_cash > 0 else 0.0
    print("\n[METRICS]")
    print(f"  总收益率: {cumulative_return:+.4%}")
    print(f"  总交易次数: {len(portfolio.trade_history)}")
    print(f"  最终 NAV: {nav:,.2f}")

    # 保存对话记录
    log_path = Path("logs/test_episode.json")
    log_path.parent.mkdir(parents=True, exist_ok=True)
    with open(log_path, "w", encoding="utf-8") as f:
        json.dump(messages, f, ensure_ascii=False, indent=2)
    print(f"\n对话记录已保存到 {log_path}")


def main() -> None:
    parser = argparse.ArgumentParser(description="独立测试 Agent — 调用外部 LLM API 运行单日交易")
    parser.add_argument("--base-url", default=None, help="OpenAI-compatible API base URL (默认从 OPENAI_BASE_URL 环境变量读取)")
    parser.add_argument("--api-key", default=None, help="API key (默认从 OPENAI_API_KEY 环境变量读取)")
    parser.add_argument("--model", default="Qwen/Qwen3-4B", help="模型名称")
    parser.add_argument("--date", default="2025-03-15", help="交易日期")
    parser.add_argument("--universe", nargs="+", default=None, help="股票池（默认使用 train_data_config）")
    parser.add_argument("--initial-cash", type=float, default=1_000_000.0, help="初始资金")
    parser.add_argument("--qlib-data-dir", default="data/qlib_data/cn_data", help="Qlib 数据目录")
    parser.add_argument("--max-turns", type=int, default=15, help="最大对话轮数")
    parser.add_argument("--enable-thinking", action="store_true", help="启用 Qwen3 的 thinking mode")

    args = parser.parse_args()

    import os
    base_url = args.base_url or os.environ.get("OPENAI_BASE_URL", "http://localhost:8000/v1")
    api_key = args.api_key or os.environ.get("OPENAI_API_KEY", "empty")

    if args.universe is None:
        from train_data_config import DEFAULT_UNIVERSE
        universe = DEFAULT_UNIVERSE[:3]
    else:
        universe = args.universe

    run_episode(
        base_url=base_url,
        api_key=api_key,
        model=args.model,
        date=args.date,
        universe=universe,
        initial_cash=args.initial_cash,
        qlib_data_dir=args.qlib_data_dir,
        max_turns=args.max_turns,
        enable_thinking=args.enable_thinking,
    )


if __name__ == "__main__":
    main()