#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Inference entry point for Track A.

This script owns the executable prediction flow:
    1. Load test scenarios.
    2. Hydrate placeholder telemetry from the competition server when needed.
    3. Load the trained model bundle produced by train.py.
    4. Ask the Qwen/OpenRouter assistant to call the trained-model tool.
    5. Checkpoint and finally write results/result.csv, debug JSON, and traces JSON.

Reusable feature extraction, prediction helpers, server utilities, and file writers
live in src/model_core.py. Training and cross-validation live in train.py.

python main.py \
  --test_path "data/Phase_2/test.json" \
  --model_bundle "results/model_v4_bundle.pkl" \
  --out "results/result.csv" \
  --debug_out "results/debug_phase2_v4.json" \
  --traces_out "results/traces.json" \
  --checkpoint_every 1 \
  --max_steps 4 \
  --max_tool_calls 8 \
  --concurrency 1 \
  --question_timeout 180
"""

from __future__ import annotations

import argparse
import os
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Optional

import httpx
from openai import OpenAI
from tqdm import tqdm

from src.model_core import (
    build_question_text,
    get_env_value,
    get_options,
    hydrate_scenario_from_server,
    load_env_file,
    load_json,
    load_model_bundle,
    result_csv_path,
    run_agent_with_trained_model,
    scenario_needs_server_data,
    write_debug,
    write_json,
    write_submission,
)

DEFAULT_TRACK_B_TEST = os.path.join("..", "Track B", "data", "Phase_2", "test.json")


def normalize_track_a_answer(answer: object) -> str:
    text = str(answer or "").strip()
    if not text:
        return ""
    labels = re.findall(r"C\d+", text)
    if not labels:
        return text
    labels = sorted(dict.fromkeys(labels), key=lambda label: int(label[1:]))
    return "|".join(labels)


def normalize_submission_rows(rows: list[dict[str, str]]) -> list[dict[str, str]]:
    out = []
    for row in rows:
        out.append(
            {
                "ID": row.get("ID", ""),
                "Track A": normalize_track_a_answer(row.get("Track A", "")),
                "Track B": str(row.get("Track B", "") or ""),
            }
        )
    return out


def write_result(path: str, rows: list[dict[str, str]]) -> None:
    write_submission(path, normalize_submission_rows(rows))


def append_track_b_blank_rows(
    rows: list[dict[str, str]], track_b_test_path: str
) -> int:
    if not track_b_test_path or not os.path.exists(track_b_test_path):
        return 0
    track_b = load_json(track_b_test_path)
    existing = {str(row.get("ID", "")) for row in rows}
    added = 0
    for i, s in enumerate(track_b, start=1):
        sid = str(s.get("scenario_id") or s.get("ID") or "").strip()
        if not sid:
            task_id = s.get("task", {}).get("id", i)
            raise ValueError(f"Missing scenario_id for Track B task id {task_id}")
        if sid in existing:
            raise ValueError(f"Track B ID already exists in output rows: {sid}")
        rows.append({"ID": sid, "Track A": "", "Track B": ""})
        existing.add(sid)
        added += 1
    return added


def main() -> None:
    load_env_file()
    parser = argparse.ArgumentParser()
    parser.add_argument("--test_path", default="data/Phase_1/test.json")
    parser.add_argument("--model_bundle", default="results/model_v4_bundle.pkl")
    parser.add_argument("--track_b_test", default=DEFAULT_TRACK_B_TEST)
    parser.add_argument("--out", default="results/result.csv")
    parser.add_argument("--debug_out", default="results/debug_v4.json")
    parser.add_argument("--traces_out", default="results/traces.json")
    parser.add_argument("--max_samples", type=int, default=None)
    parser.add_argument("--use_server_data", action="store_true")
    parser.add_argument("--no_auto_server_data", action="store_true")
    parser.add_argument("--server_url", default="https://124.71.227.61/no")
    parser.add_argument("--env_path", default=None)
    parser.add_argument("--auth_token_env", default="AUTH_TOKEN")
    parser.add_argument("--timeout", type=float, default=30.0)
    parser.add_argument("--verify_ssl", action="store_true")
    parser.add_argument("--try_scenario_endpoint", action="store_true")
    parser.add_argument("--checkpoint_every", type=int, default=1)
    parser.add_argument("--no_progress", action="store_true")
    parser.add_argument("--no_agent", action="store_true")
    parser.add_argument("--max_steps", type=int, default=4)
    parser.add_argument("--max_tool_calls", type=int, default=8)
    parser.add_argument("--concurrency", type=int, default=1)
    parser.add_argument(
        "--model_url",
        default=os.getenv("OPENROUTER_URL")
        or os.getenv("OPENAI_BASE_URL")
        or "https://openrouter.ai/api/v1",
    )
    parser.add_argument(
        "--model_name",
        default=os.getenv("OPENROUTER_MODEL")
        or os.getenv("OPENAI_MODEL")
        or "qwen/qwen3.5-35b-a3b",
    )
    parser.add_argument(
        "--api_key_env",
        default="OPENROUTER_API_KEY,AGENT_API_KEY,OPENAI_API_KEY",
    )
    parser.add_argument("--agent_timeout", type=float, default=60.0)
    parser.add_argument("--llm_timeout", type=float, default=None)
    parser.add_argument("--question_timeout", type=float, default=180.0)
    parser.add_argument("--max_output_tokens", type=int, default=900)
    parser.add_argument("--history_chars", type=int, default=24000)
    parser.add_argument("--observation_chars", type=int, default=20000)
    parser.add_argument("--temperature", type=float, default=0.0)
    args = parser.parse_args()
    args.out = result_csv_path(args.out)
    llm_timeout = args.llm_timeout if args.llm_timeout is not None else args.agent_timeout

    test = load_json(args.test_path)
    if args.max_samples is not None:
        test = test[: max(0, args.max_samples)]
    print(f"Loaded test={len(test)}")

    loaded_env = load_env_file(args.env_path)

    auto_server_data = not args.no_auto_server_data and any(
        scenario_needs_server_data(s) for s in test
    )
    use_server_data = args.use_server_data or auto_server_data
    if use_server_data:
        print(
            f"Server data mode enabled: {args.server_url}"
            + (f" (env: {loaded_env})" if loaded_env else "")
        )

    if not os.path.exists(args.model_bundle):
        raise FileNotFoundError(
            f"Model bundle not found: {args.model_bundle}. Run train.py first."
        )
    print(f"Loading trained model bundle: {args.model_bundle}")
    tmodel, smodel, model_metadata = load_model_bundle(args.model_bundle)
    if model_metadata:
        print(f"Model metadata: {model_metadata}")

    use_agent = not args.no_agent
    llm_client: Optional[OpenAI] = None
    if use_agent:
        api_key = get_env_value(args.api_key_env)
        if api_key:
            llm_client = OpenAI(
                base_url=args.model_url,
                api_key=api_key,
                http_client=httpx.Client(verify=args.verify_ssl),
                timeout=llm_timeout,
            )
            print(f"Agent mode enabled: {args.model_name}")
        else:
            print(
                f"Warning: {args.api_key_env} is not set; using direct trained-model fallback."
            )

    row_entries, debug_entries, trace_entries = [], [], []
    print("Predicting Track A test...")
    headers = {"Content-Type": "application/json"}
    token = os.environ.get(args.auth_token_env, "").strip()
    if token:
        headers["Authorization"] = f"Bearer {token}"
        headers["X-API-Token"] = token
    if use_server_data and not token:
        print(f"Warning: {args.auth_token_env} is not set; server may reject requests.")

    def ordered_rows() -> list[dict[str, str]]:
        return [row for _, row in sorted(row_entries, key=lambda x: x[0])]

    def ordered_debug() -> list[dict]:
        return [row for _, row in sorted(debug_entries, key=lambda x: x[0])]

    def ordered_traces() -> list[dict]:
        return [row for _, row in sorted(trace_entries, key=lambda x: x[0])]

    def checkpoint() -> None:
        write_result(args.out, ordered_rows())
        write_debug(args.debug_out, ordered_debug())
        write_json(args.traces_out, ordered_traces())

    def solve_one(i: int, s: dict) -> tuple[int, dict, dict, dict, str]:
        sid = s.get("scenario_id") or s.get("ID") or f"test_{i}"
        start_time = time.perf_counter()
        tool_calls: List[str] = []
        pred_s = s
        used_server_data = False
        try:
            if use_server_data and (
                args.use_server_data or scenario_needs_server_data(s)
            ):
                with httpx.Client(
                    headers=headers,
                    timeout=args.timeout,
                    verify=args.verify_ssl,
                    follow_redirects=True,
                ) as client_ctx:
                    pred_s, tool_calls = hydrate_scenario_from_server(
                        s,
                        client_ctx,
                        args.server_url,
                        try_scenario_endpoint=args.try_scenario_endpoint,
                    )
                used_server_data = bool(tool_calls)

            labels, dbg, completion, agent_tool_calls, agent_used = (
                run_agent_with_trained_model(
                    llm_client,
                    args.model_name,
                    tmodel,
                    smodel,
                    pred_s,
                    timeout=llm_timeout,
                    max_steps=args.max_steps,
                    max_tool_calls=args.max_tool_calls,
                    temperature=args.temperature,
                    max_output_tokens=args.max_output_tokens,
                    history_chars=args.history_chars,
                    observation_chars=args.observation_chars,
                    question_timeout=args.question_timeout,
                )
            )
            tool_calls.extend(agent_tool_calls)
            pred = "|".join(labels)
            elapsed = round(time.perf_counter() - start_time, 3)
            row = {"ID": sid, "Track A": pred, "Track B": ""}
            debug_row = {
                "scenario_id": sid,
                "prediction": pred,
                "debug": dbg,
                "options": get_options(pred_s),
                "used_server_data": used_server_data,
                "still_needs_server_data": scenario_needs_server_data(pred_s),
                "agent_used": agent_used,
            }
            trace_row = {
                "scenario_id": sid,
                "question": build_question_text(pred_s),
                "completion": completion,
                "prediction": labels,
                "ground_truth": pred_s.get("answer", "To be determined"),
                "score": 0.0,
                "execution_time_seconds": elapsed,
                "tool_calls": "\n".join(tool_calls),
                "boxed": f"\\boxed{{{pred}}}",
            }
            prob = dbg.get("template_prob", 0.0)
            try:
                prob_text = f"{float(prob):.3f}"
            except Exception:
                prob_text = "nan"
            msg = (
                f"[A {i}/{len(test)}] {sid} -> {pred} "
                f"({dbg.get('template')} {prob_text})"
            )
            return i, row, debug_row, trace_row, msg
        except Exception as exc:
            elapsed = round(time.perf_counter() - start_time, 3)
            row = {"ID": sid, "Track A": "", "Track B": ""}
            debug_row = {
                "scenario_id": sid,
                "prediction": "",
                "debug": {
                    "runner_exception": f"{type(exc).__name__}: {exc}",
                    "agent_used": False,
                },
                "options": get_options(pred_s),
                "used_server_data": used_server_data,
                "still_needs_server_data": scenario_needs_server_data(pred_s),
                "agent_used": False,
            }
            trace_row = {
                "scenario_id": sid,
                "question": build_question_text(pred_s),
                "completion": "",
                "prediction": [],
                "ground_truth": pred_s.get("answer", "To be determined"),
                "score": 0.0,
                "execution_time_seconds": elapsed,
                "tool_calls": "\n".join(tool_calls),
                "boxed": "\\boxed{}",
                "messages": [
                    {
                        "action": "runner_exception",
                        "observation": f"{type(exc).__name__}: {exc}",
                    }
                ],
            }
            return i, row, debug_row, trace_row, f"[A {i}/{len(test)}] {sid} -> ERROR: {exc}"

    max_workers = max(1, int(args.concurrency))
    if max_workers == 1:
        iterator = enumerate(test, start=1)
        if not args.no_progress:
            iterator = tqdm(iterator, total=len(test), desc="Predicting Track A")
        for i, s in iterator:
            result_i, row, debug_row, trace_row, msg = solve_one(i, s)
            row_entries.append((result_i, row))
            debug_entries.append((result_i, debug_row))
            trace_entries.append((result_i, trace_row))
            if args.checkpoint_every and result_i % args.checkpoint_every == 0:
                checkpoint()
            if result_i <= 5 or result_i % 50 == 0:
                if args.no_progress:
                    print(msg)
                else:
                    tqdm.write(msg)
    else:
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_index = {
                executor.submit(solve_one, i, s): i for i, s in enumerate(test, start=1)
            }
            iterator = as_completed(future_to_index)
            if not args.no_progress:
                iterator = tqdm(iterator, total=len(test), desc="Predicting Track A")
            completed = 0
            for future in iterator:
                completed += 1
                result_i, row, debug_row, trace_row, msg = future.result()
                row_entries.append((result_i, row))
                debug_entries.append((result_i, debug_row))
                trace_entries.append((result_i, trace_row))
                if args.checkpoint_every and completed % args.checkpoint_every == 0:
                    checkpoint()
                if args.no_progress:
                    print(f"[{completed}/{len(test)}] {msg}")
                else:
                    tqdm.write(f"[{completed}/{len(test)}] {msg}")

    rows = ordered_rows()
    added_track_b = append_track_b_blank_rows(rows, args.track_b_test)
    if added_track_b:
        print(f"Added Track B blank ID rows: {added_track_b}")

    write_result(args.out, rows)
    write_debug(args.debug_out, ordered_debug())
    write_json(args.traces_out, ordered_traces())
    print(f"Saved result: {args.out}")
    print(f"Saved debug: {args.debug_out}")
    print(f"Saved traces: {args.traces_out}")


if __name__ == "__main__":
    main()