"""Test a deployed model on Alibaba Cloud with U-MATH validation samples.

Multi-turn agentic rollout with python_interpreter tool
(reuses CalcXRolloutProcessor from functions/rollout),
then evaluates with the full reward pipeline (format 0.2 + accuracy 0.8).

Usage:
    # Debug mode: 5 samples, full trajectory log
    python test/test_deployed_umath.py
    python test/test_deployed_umath.py --loglevel INFO --num 10

    # Full eval mode: all samples, export trajectory to jsonl
    python test/test_deployed_umath.py --all
    python test/test_deployed_umath.py --all --output results/my_run.jsonl
"""
import argparse
import asyncio
import json
import logging
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
os.environ.setdefault("NO_PROXY", "localhost,127.0.0.1,::1")

import httpx

from functions.reward.math_verify import extract_answer_by_format

from dashscope.finetune.reinforcement import RolloutInput
from dashscope.finetune.reinforcement.component.data.base_data_model import (
    ModelResource, TaskStatus,
)

MODEL_CODE = "<your-deployed-model-code>"
UMATH_TEST_FILE = os.path.join(os.path.dirname(__file__), "..", "data", "umath_test.jsonl")
LLM_CONFIG_PATH = os.path.join(os.path.dirname(__file__), "..", "llm_api_config.json")

JUDGE_PROMPT = """\
You are a math grading assistant. Given the original problem, compare the \
student's answer with the reference answer and decide how much of it is \
mathematically correct.

**Original problem:**
{question}

**Reference answer:**
{golden_answer}

**Student's full response (may contain reasoning):**
{model_response}

Scoring rules:
  - 1.0  : fully correct (mathematically equivalent to the reference)
  - 0.0  : fully incorrect or no answer
  - 0 < score < 1.0 : partial credit, ONLY when the reference contains \
multiple distinct parts (e.g. several sub-questions, multiple items in a \
list, a function plus its domain, or a tuple of independent values). \
Set score = (number of parts the student got right) / (total parts).
Single-part questions only get 0.0 or 1.0.

First think briefly about whether the answers match, then output your decision.
Reply with ONLY a single JSON object, no markdown fences, no extra prose. The \
"reason" field MUST come before the "score" field:
{{"reason": "<one short sentence>", "score": <number between 0.0 and 1.0>}}"""

logger = logging.getLogger("test_deployed_umath")

# Third-party loggers to suppress — only our own logger stays at the user-chosen level
_NOISY_LOGGERS = [
    "httpcore", "httpx", "openai", "openai._base_client",
    "mcp", "mcp.client", "mcp.client.sse", "mcp.server",
    "mcp.server.lowlevel.server", "urllib3", "asyncio",
    "uvicorn", "uvicorn.access", "uvicorn.error",
    "langchain_openai", "langchain", "langgraph",
    "functions.rollout.math_rollout",
]

# Collect request IDs and model from each model call via a logging filter on openai
_request_ids: list[str] = []
_request_models: list[str] = []


class _RequestIdCapture(logging.Filter):
    """Intercept openai._base_client DEBUG lines to capture request_id and model."""
    def filter(self, record: logging.LogRecord) -> bool:
        msg = record.getMessage()
        if "request_id:" in msg:
            rid = msg.split("request_id:")[-1].strip()
            _request_ids.append(rid)
        if "'model':" in msg and "json_data" in msg:
            try:
                start = msg.index("'model':") + len("'model':")
                segment = msg[start:start+80].strip().strip("'\"").split("'")[0].split('"')[0].strip(", ")
                _request_models.append(segment)
            except Exception:
                pass
        return False  # suppress all openai logs from output


# ============================================================================ #
#                                CONFIG / DATA                                  #
# ============================================================================ #

def load_config():
    with open(LLM_CONFIG_PATH, encoding="utf-8") as f:
        cfg = json.load(f)
    api_key = os.environ.get("DASHSCOPE_API_KEY") or cfg["key"]
    base_url = cfg["base_url"]
    judge_model = cfg.get("supported_models", ["qwen-plus"])[-1]
    return api_key, base_url, judge_model


def load_samples(n: int | None = None):
    samples = []
    with open(UMATH_TEST_FILE, encoding="utf-8") as f:
        for i, line in enumerate(f):
            if n is not None and i >= n:
                break
            samples.append(json.loads(line))
    return samples


# ============================================================================ #
#                              REWARD EVALUATION                                #
# ============================================================================ #

def llm_judge(model_response: str, golden_answer: str, question: str, config: dict) -> dict:
    prompt = JUDGE_PROMPT.format(
        question=question,
        golden_answer=golden_answer,
        model_response=model_response[-3000:],
    )
    try:
        resp = httpx.post(
            f"{config['base_url']}/chat/completions",
            headers={
                "Authorization": f"Bearer {config['api_key']}",
                "Content-Type": "application/json",
            },
            json={
                "model": config["judge_model"],
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.0,
                "max_tokens": 200,
                "response_format": {"type": "json_object"},
            },
            timeout=60.0,
        )
        resp.raise_for_status()
        text = resp.json()["choices"][0]["message"]["content"].strip()
        logger.debug(f"Judge raw response: {text}")
        data = json.loads(text[text.find("{"):text.rfind("}") + 1])
        score = max(0.0, min(1.0, float(data.get("score", 0.0))))
        reason = data.get("reason", "")
        return {"score": score, "reason": reason, "raw": text}
    except Exception as e:
        logger.error(f"Judge call failed: {e}")
        return {"score": 0.0, "reason": f"error: {e}", "raw": None}


def evaluate_reward(response: str, golden: str, question: str, judge_config: dict) -> dict:
    """format(0.2) + accuracy(0.8) with LLM judge."""
    extracted = extract_answer_by_format(response)
    format_score = 1.0 if extracted is not None else -1.0

    if extracted is None:
        accuracy_score = 0.0
        judge_result = {"score": 0.0, "reason": "no \\boxed{} found, skipped judge", "raw": None}
    else:
        judge_result = llm_judge(response, golden, question, judge_config)
        accuracy_score = judge_result["score"]

    total = 0.2 * format_score + 0.8 * accuracy_score

    return {
        "format_score": format_score,
        "accuracy_score": accuracy_score,
        "total_reward": total,
        "extracted_answer": extracted,
        "judge_reason": judge_result["reason"],
        "judge_raw": judge_result["raw"],
    }


# ============================================================================ #
#                         MULTI-TURN ROLLOUT (with tools)                       #
# ============================================================================ #

async def run_rollout(sample: dict, processor, api_key: str, base_url: str, model: str,
                      track_requests: bool = True) -> tuple[str, dict]:
    """Multi-turn agentic rollout via CalcXRolloutProcessor."""
    rollout_input = RolloutInput(
        messages=sample["messages"],
        rollout_extra=sample.get("rollout_extra"),
        model_resource=ModelResource(
            model_name=model,
            base_url=base_url,
            api_key=api_key,
        ),
        sampling_params={
            "temperature": 0.6,
            "max_tokens": 8192,
            "max_turns": 10,
        },
    )

    if track_requests:
        _request_ids.clear()
        _request_models.clear()

    t0 = time.time()
    result = await processor.process(rollout_input)
    latency = time.time() - t0

    request_ids = list(_request_ids) if track_requests else []
    request_models = list(_request_models) if track_requests else []

    messages = result.agent_output.messages if result.agent_output else []
    final_response = ""
    if messages:
        last = messages[-1]
        final_response = last.get("content", "") if isinstance(last, dict) else str(last)

    tool_call_count = 0
    for msg in messages:
        if isinstance(msg, dict) and msg.get("role") == "assistant":
            tc = msg.get("tool_calls")
            if tc:
                tool_call_count += len(tc)

    rollout_metrics = {}
    if result.agent_output and result.agent_output.rollout_metrics:
        rollout_metrics = result.agent_output.rollout_metrics

    meta = {
        "latency": latency,
        "tool_calls": rollout_metrics.get("tool_call_count", tool_call_count),
        "tool_total_time": rollout_metrics.get("tool_call_total_time", 0),
        "turns": len([m for m in messages if isinstance(m, dict) and m.get("role") == "assistant"]),
        "status": result.status.value if result.status else "unknown",
        "error": result.error,
        "messages": messages,
        "request_ids": request_ids,
        "request_models": request_models,
    }
    return final_response, meta


def init_rollout_processor():
    from functions.rollout.math_rollout import CalcXRolloutProcessor
    logger.info("Initializing CalcXRolloutProcessor (starting MCP server)...")
    processor = CalcXRolloutProcessor()
    logger.info("CalcXRolloutProcessor ready.")
    return processor


# ============================================================================ #
#                          TRAJECTORY PRETTY PRINT                              #
# ============================================================================ #

def print_trajectory(messages: list, request_ids: list | None = None):
    """Print multi-turn conversation in a readable format."""
    logger.debug("")
    logger.debug("┌─── Trajectory ───────────────────────────────────────────────")

    # Map request_ids to assistant turns (one request_id per model call)
    rid_iter = iter(request_ids or [])

    for j, msg in enumerate(messages):
        if not isinstance(msg, dict):
            continue
        role = msg.get("role", "?")
        content = msg.get("content", "")
        tool_calls = msg.get("tool_calls")

        if role == "system":
            logger.debug(f"│ [{j}] 🔧 SYSTEM:")
            for line in content.split("\n"):
                logger.debug(f"│     {line}")
        elif role == "user":
            logger.debug(f"│ [{j}] 👤 USER:")
            for line in content.split("\n"):
                logger.debug(f"│     {line}")
        elif role == "assistant" and tool_calls:
            rid = next(rid_iter, None)
            rid_str = f"  request_id={rid}" if rid else ""
            logger.debug(f"│ [{j}] 🤖 ASSISTANT (tool_call){rid_str}:")
            if content:
                for line in content.split("\n"):
                    logger.debug(f"│     {line}")
            for tc in tool_calls:
                fn = tc.get("function", {})
                name = fn.get("name", "?")
                args = fn.get("arguments", "")
                logger.debug(f"│     ↳ {name}(")
                for line in args.split("\n"):
                    logger.debug(f"│       {line}")
                logger.debug(f"│     )")
        elif role == "assistant":
            rid = next(rid_iter, None)
            rid_str = f"  request_id={rid}" if rid else ""
            logger.debug(f"│ [{j}] 🤖 ASSISTANT{rid_str}:")
            for line in content.split("\n"):
                logger.debug(f"│     {line}")
        elif role == "tool":
            tool_name = msg.get("name", "?")
            if isinstance(content, list):
                text = "\n".join(
                    item.get("text", "") for item in content if isinstance(item, dict)
                )
            else:
                text = str(content)
            logger.debug(f"│ [{j}] 🔨 TOOL ({tool_name}):")
            for line in text.split("\n"):
                logger.debug(f"│     {line}")
        else:
            logger.debug(f"│ [{j}] {role}: {content}")

    logger.debug("└───────────────────────────────────────────────────────────────")
    logger.debug("")


# ============================================================================ #
#                                   MAIN                                        #
# ============================================================================ #

def print_summary(results: list, model: str, _print=print):
    n = len(results)
    if not n:
        _print("No results.")
        return

    avg_reward = sum(r["total_reward"] for r in results) / n
    avg_accuracy = sum(r["accuracy_score"] for r in results) / n
    format_rate = sum(1 for r in results if r["format_score"] > 0) / n
    avg_latency = sum(r["latency"] for r in results) / n
    avg_tool_calls = sum(r["tool_calls"] for r in results) / n
    avg_turns = sum(r["turns"] for r in results) / n

    _print(f"\n{'='*80}")
    _print(f"SUMMARY  (model={model}, n={n})")
    _print(f"{'='*80}")
    _print(f"Avg total reward:  {avg_reward:.3f}")
    _print(f"Avg accuracy:      {avg_accuracy:.3f}")
    _print(f"Format rate:       {format_rate:.1%}")
    _print(f"Avg latency:       {avg_latency:.2f}s")
    _print(f"Avg tool calls:    {avg_tool_calls:.1f}")
    _print(f"Avg turns:         {avg_turns:.1f}")
    _print(f"\nPer-sample:")
    for r in results:
        _print(f"  [{r['index']}] reward={r['total_reward']:.3f} "
               f"acc={r['accuracy_score']:.3f} tools={r['tool_calls']} "
               f"turns={r['turns']} latency={r['latency']:.1f}s "
               f"| {r['subject']}")


async def async_main():
    parser = argparse.ArgumentParser(description="Test deployed umath model")
    parser.add_argument("--loglevel", default="DEBUG",
                        help="Log level: DEBUG (full output) / INFO / WARNING")
    parser.add_argument("--num", type=int, default=5, help="Number of samples (ignored if --all)")
    parser.add_argument("--model", default=MODEL_CODE, help="Deployed model code")
    parser.add_argument("--all", action="store_true",
                        help="Evaluate all samples, export trajectory to jsonl")
    parser.add_argument("--output", type=str, default=None,
                        help="Output jsonl path (default: test/results/<model>_<timestamp>.jsonl)")
    parser.add_argument("--resume", type=str, default=None,
                        help="Resume from an existing jsonl file, skip completed samples and append")
    parser.add_argument("--concurrency", type=int, default=1,
                        help="Number of concurrent rollouts (default 1; use with --all)")
    parser.add_argument("--timeout", type=int, default=180,
                        help="Per-sample timeout in seconds (default 180)")
    args = parser.parse_args()

    # --resume: reuse the file, but don't force --all if user wants debug
    if args.resume:
        args.output = args.resume
        if not args.all:
            # resume + debug mode: load all samples but show full trajectory
            pass
        else:
            args.all = True

    # --all mode: force WARNING to keep screen clean, default to all samples
    if args.all:
        level = logging.WARNING
        num_samples = None
    else:
        level = getattr(logging, args.loglevel.upper(), logging.DEBUG)
        num_samples = None if args.resume else args.num

    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
        datefmt="%H:%M:%S",
    )
    for name in _NOISY_LOGGERS:
        logging.getLogger(name).setLevel(logging.WARNING)
    _openai_logger = logging.getLogger("openai._base_client")
    _openai_logger.setLevel(logging.DEBUG)
    _openai_logger.addFilter(_RequestIdCapture())

    api_key, base_url, judge_model = load_config()
    judge_config = {"api_key": api_key, "base_url": base_url, "judge_model": judge_model}
    samples = load_samples(num_samples)

    # Determine output path for --all mode
    output_path = None
    if args.all:
        if args.output:
            output_path = args.output
        else:
            results_dir = os.path.join(os.path.dirname(__file__), "results")
            os.makedirs(results_dir, exist_ok=True)
            ts = time.strftime("%Y%m%d_%H%M%S")
            output_path = os.path.join(results_dir, f"{args.model}_{ts}.jsonl")

    # In --all mode, redirect fd 1 & 2 to devnull so the MCP subprocess
    # (uvicorn access logs) and rich console output are silenced.
    # We keep a clean fd for our own progress output.
    _saved_stdout_fd = None
    _progress_fd = None
    if args.all:
        _saved_stdout_fd = os.dup(1)
        _progress_fd = os.fdopen(_saved_stdout_fd, "w", closefd=False)
        _devnull_fd = os.open(os.devnull, os.O_WRONLY)
        os.dup2(_devnull_fd, 1)
        os.dup2(_devnull_fd, 2)
        os.close(_devnull_fd)

    def _print(*a, **kw):
        """Print to progress fd (bypasses devnull redirect in --all mode)."""
        if _progress_fd:
            kw["file"] = _progress_fd
            kw.setdefault("flush", True)
        builtins_print(*a, **kw)

    builtins_print = print

    _print(f"Model: {args.model}")
    _print(f"Judge model: {judge_model}")
    _print(f"Samples: {len(samples)}")
    if args.concurrency > 1:
        _print(f"Concurrency: {args.concurrency}")
    if args.resume:
        _print(f"Resume from: {args.resume}")
    elif output_path:
        _print(f"Output: {output_path}")
    _print("=" * 80)

    processor = init_rollout_processor()

    from tqdm import tqdm

    # Resume: load completed indices from existing file
    done_indices: set[int] = set()
    results = []
    if args.resume and os.path.exists(args.resume):
        with open(args.resume, encoding="utf-8") as f:
            for line in f:
                rec = json.loads(line)
                done_indices.add(rec["index"])
                results.append(rec)
        _print(f"Resuming: {len(done_indices)} samples already done, skipping them")

    outfile = open(output_path, "a" if args.resume else "w", encoding="utf-8") if output_path else None
    # tqdm progress bar: only in --all mode, write to our saved fd
    pbar = tqdm(total=len(samples), initial=len(done_indices),
                file=_progress_fd or sys.stderr,
                desc="Evaluating", unit="sample",
                bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]"
                ) if args.all else None
    running_acc = [r["accuracy_score"] for r in results]
    concurrent = args.concurrency > 1 and args.all
    _outfile_lock = asyncio.Lock()

    sample_timeout = args.timeout

    async def _run_with_timeout(sample, track_requests):
        """Run rollout with per-sample timeout."""
        try:
            return await asyncio.wait_for(
                run_rollout(sample, processor, api_key, base_url, args.model,
                            track_requests=track_requests),
                timeout=sample_timeout,
            )
        except asyncio.TimeoutError:
            return "", {
                "latency": sample_timeout,
                "tool_calls": 0,
                "turns": 0,
                "status": "timeout",
                "error": f"timeout after {sample_timeout}s",
                "messages": [],
                "request_ids": [],
                "request_models": [],
            }

    async def _process_one(i: int, sample: dict, sem: asyncio.Semaphore | None = None):
        """Process a single sample (rollout + reward). Works both sequential and concurrent."""
        user_msg = sample["messages"][0]["content"]
        golden = sample.get("rollout_extra", {}).get("solution", "N/A")
        subject = sample.get("rollout_extra", {}).get("subject", "unknown")

        if sem:
            async with sem:
                response, meta = await _run_with_timeout(sample, track_requests=False)
        else:
            response, meta = await _run_with_timeout(sample, track_requests=True)

        reward = evaluate_reward(response, golden, user_msg, judge_config)

        result_record = {
            "index": i,
            "subject": subject,
            "question": user_msg,
            "golden": golden,
            "extracted_answer": reward["extracted_answer"],
            "response": response,
            "messages": meta.get("messages", []),
            "request_ids": meta.get("request_ids", []),
            "latency": meta["latency"],
            "tool_calls": meta["tool_calls"],
            "turns": meta["turns"],
            "status": meta.get("status"),
            "error": meta.get("error"),
            **reward,
        }

        # Thread-safe output & progress update
        async with _outfile_lock:
            results.append(result_record)
            if outfile:
                outfile.write(json.dumps(result_record, ensure_ascii=False) + "\n")
                outfile.flush()
            if pbar:
                running_acc.append(reward["accuracy_score"])
                avg_acc = sum(running_acc) / len(running_acc)
                status_icon = "+" if reward["total_reward"] >= 0.5 else "-"
                pbar.set_postfix_str(f"acc={avg_acc:.3f} | {status_icon} {subject}")
                pbar.update(1)

        return result_record

    try:
        if concurrent:
            # Concurrent mode: fan out with semaphore
            sem = asyncio.Semaphore(args.concurrency)
            tasks = []
            for i, sample in enumerate(samples):
                if i in done_indices:
                    continue
                tasks.append(_process_one(i, sample, sem))
            await asyncio.gather(*tasks, return_exceptions=True)
        else:
            # Sequential mode: full debug output
            for i, sample in enumerate(samples):
                if i in done_indices:
                    continue
                user_msg = sample["messages"][0]["content"]
                golden = sample.get("rollout_extra", {}).get("solution", "N/A")
                subject = sample.get("rollout_extra", {}).get("subject", "unknown")

                if not args.all:
                    _print(f"\n{'='*80}")
                    _print(f"[Sample {i+1}/{len(samples)}] subject={subject}")
                    _print(f"{'='*80}")
                    logger.debug(f"Question: {user_msg}")
                    logger.debug(f"Golden:   {golden}")

                response, meta = await _run_with_timeout(sample, track_requests=True)

                if not args.all:
                    logger.info(f"Latency: {meta['latency']:.2f}s | Tool calls: {meta['tool_calls']} | "
                                f"Turns: {meta['turns']} | Status: {meta.get('status', '?')}")
                    logger.info(f"Models called: {meta.get('request_models', [])}")
                    if meta.get("error"):
                        logger.error(f"Rollout error: {meta['error']}")
                    print_trajectory(meta.get("messages", []), meta.get("request_ids"))

                reward = evaluate_reward(response, golden, user_msg, judge_config)

                if not args.all:
                    logger.info(f"Reward: total={reward['total_reward']:.3f} "
                                f"(format={reward['format_score']:.1f}, accuracy={reward['accuracy_score']:.3f})")
                    logger.info(f"Golden:    {golden}")
                    logger.info(f"Extracted: {reward['extracted_answer']}")
                    logger.info(f"Judge reason: {reward['judge_reason']}")
                    logger.debug(f"Judge raw: {reward['judge_raw']}")
                else:
                    running_acc.append(reward["accuracy_score"])
                    avg_acc = sum(running_acc) / len(running_acc)
                    status_icon = "+" if reward["total_reward"] >= 0.5 else "-"
                    pbar.set_postfix_str(f"acc={avg_acc:.3f} | {status_icon} {subject}")
                    pbar.update(1)

                result_record = {
                    "index": i,
                    "subject": subject,
                    "question": user_msg,
                    "golden": golden,
                    "extracted_answer": reward["extracted_answer"],
                    "response": response,
                    "messages": meta.get("messages", []),
                    "request_ids": meta.get("request_ids", []),
                    "latency": meta["latency"],
                    "tool_calls": meta["tool_calls"],
                    "turns": meta["turns"],
                    "status": meta.get("status"),
                    "error": meta.get("error"),
                    **reward,
                }
                results.append(result_record)
                if outfile:
                    outfile.write(json.dumps(result_record, ensure_ascii=False) + "\n")
                    outfile.flush()

    finally:
        if pbar:
            pbar.close()
        if outfile:
            outfile.close()

    print_summary(results, args.model, _print)

    if output_path:
        _print(f"\nTrajectory exported to: {output_path}")


if __name__ == "__main__":
    asyncio.run(async_main())
