#!/usr/bin/env python3
"""
SDR Benchmark Runner + Evaluator using Claude Haiku.

Replaces the DeepResearchGym eval_quality_async.py OpenAI dependency
with Anthropic Claude Haiku (claude-haiku-4-5-20251001).

Usage:
    cd /root/sdr_repo
    uv run python /root/Strategy/research/sdr_benchmark_with_eval.py --n 1 --eval-only
    uv run python /root/Strategy/research/sdr_benchmark_with_eval.py --n 3

Steps:
    1. Generate SDR reports (using configured Groq model)
    2. Evaluate reports using Claude Haiku as judge
    3. Save results to markdown
"""

import asyncio
import json
import sys
import os
import argparse
from pathlib import Path
from datetime import datetime

# Add SDR to path
sys.path.insert(0, "/root/sdr_repo")
sys.path.insert(0, "/root/sdr_repo/src")

# --- Configuration ---

SCIENTIFIC_QUERIES = [
    {"id": "sci_001", "query": "What are the state-of-the-art methods for operator splitting in machine learning optimization?"},
    {"id": "sci_002", "query": "How does FlashAttention improve transformer training efficiency compared to standard attention?"},
    {"id": "sci_003", "query": "What is the current evidence on neural operators for solving partial differential equations?"},
    {"id": "sci_004", "query": "How do large language models perform on scientific reasoning benchmarks like GPQA and SciCode?"},
    {"id": "sci_005", "query": "What are the most effective approaches for few-shot learning in low-resource scientific domains?"},
    {"id": "sci_006", "query": "How does mechanistic interpretability help understand transformer circuits and features?"},
    {"id": "sci_007", "query": "What are the convergence guarantees for proximal gradient methods with regularization?"},
    {"id": "sci_008", "query": "How do multi-agent LLM systems compare to single-agent systems on complex research tasks?"},
    {"id": "sci_009", "query": "What is the role of synthetic data in improving scientific AI systems?"},
    {"id": "sci_010", "query": "How effective are retrieval-augmented generation systems for scientific literature synthesis?"},
]

OUTPUT_DIR = Path("/root/Strategy/research/sdr_reports/SDR")

EVAL_CRITERIA = [
    {
        "name": "Clarity",
        "description": "Assess how clearly, rigorously, and analytically distinct the answer is. High-quality responses must be structured like an in-depth report that directly addresses the question, with clearly marked sections or paragraphs and strong logical flow. Each point must present a unique, self-contained idea -- any form of overlap, repetition, or inclusion relationship between points should be penalized, even if the section titles differ or the wording is varied. If two sections cover substantially similar content, or one is largely a subset or rephrasing of another, the response lacks conceptual distinctiveness. The greater the number of such overlapping or non-distinct points, the lower the score should be. Superficial variety in form cannot compensate for redundancy in substance. The text must avoid ambiguity, redundancy, and conversational filler. Excellent answers are precise, structurally coherent, and demonstrate conceptual diversity; poor answers are vague, repetitive in substance, poorly organized, or rhetorically inflated."
    },
    {
        "name": "Depth",
        "description": "Assess the comprehensiveness and analytical depth of the report. Excellent reports demonstrate critical thinking, nuanced analysis, and/or synthesis of information. Simply elaborating on surface-level facts is not sufficient. Word count alone does not equate to depth. Poor reports are shallow or omit key dimensions of the topic. If the answer lists multiple subtopics but does not explain them with examples, nuance, or source grounding, it should not exceed 5."
    },
    {
        "name": "Balance",
        "description": "Evaluate the fairness and objectivity of the answer. Excellent reports present multiple perspectives fairly and impartially, especially for controversial or multi-faceted topics. Poor reports show clear bias, favor one side without justification, or ignore opposing views."
    },
    {
        "name": "Breadth",
        "description": "Evaluate how many distinct and relevant subtopics, perspectives, or contexts are covered. Excellent reports provide a wide-ranging yet focused exploration -- e.g., including legal, historical, cultural, or ethical angles where appropriate. Simply presenting both sides of a binary debate is not sufficient for a high score."
    },
    {
        "name": "Support",
        "description": "Evaluate the extent to which all key claims are substantiated by specific, identifiable, and credible evidence. Providing URLs in the report is the most basic requirement. If no section (such as references or sources) provides source URLs, the score should be zero. Having URLs only meets the minimum standard and does not merit a high score. Every factual claim must be attributed to a verifiable source. Quantitative claims require precise, contextualized data. Qualitative claims must be supported by concrete examples, not hypotheticals. Sources must be cited explicitly and be traceable."
    },
    {
        "name": "Insightfulness",
        "description": "Assess how insightful the answer is. Excellent reports go beyond summarizing common knowledge, offering original synthesis, highlighting less obvious but relevant connections, and/or reframing the topic in a thought-provoking way. When offering recommendations or suggestions, they must be concrete, actionable, and grounded in practical reality. Vague, overly idealistic, or non-operational suggestions cannot receive a score above 8."
    },
]


def get_anthropic_key() -> str:
    """Get Anthropic API key from environment or Claude Code credentials."""
    key = os.environ.get("ANTHROPIC_API_KEY")
    if key:
        return key

    creds_path = Path.home() / ".claude" / ".credentials.json"
    if creds_path.exists():
        with open(creds_path) as f:
            data = json.load(f)
        oauth = data.get("claudeAiOauth", {})
        token = oauth.get("accessToken")
        if token:
            return token

    raise RuntimeError(
        "No Anthropic API key found. Set ANTHROPIC_API_KEY or ensure Claude Code credentials exist."
    )


# --- Report Generation ---

async def run_single_query(graph, query_id: str, query_text: str, output_dir: Path) -> dict:
    """Run SDR on a single query and save output."""
    from langchain_core.messages import HumanMessage
    from open_deep_research.configuration import Config

    print(f"[{datetime.now():%H:%M:%S}] Running query {query_id}: {query_text[:80]}...")

    q_file = output_dir / f"{query_id}.q"
    q_file.write_text(query_text, encoding="utf-8")

    config = Config()

    try:
        result = await graph.ainvoke(
            {"messages": [HumanMessage(content=query_text)]},
            config={
                "configurable": {
                    **config.model_dump(),
                    "allow_clarification": False,
                }
            }
        )

        report = ""
        for msg in reversed(result.get("messages", [])):
            if hasattr(msg, "content") and isinstance(msg.content, str) and len(msg.content) > 500:
                report = msg.content
                break

        if not report:
            report = result.get("final_report", "No report generated")

        a_file = output_dir / f"{query_id}.a"
        a_file.write_text(report, encoding="utf-8")

        print(f"[{datetime.now():%H:%M:%S}] {query_id}: {len(report)} chars saved")
        return {"id": query_id, "status": "ok", "chars": len(report)}

    except Exception as e:
        error_msg = f"ERROR: {type(e).__name__}: {e}"
        print(f"[{datetime.now():%H:%M:%S}] {query_id}: FAILED - {error_msg}")
        a_file = output_dir / f"{query_id}.a"
        a_file.write_text(error_msg, encoding="utf-8")
        return {"id": query_id, "status": "error", "error": str(e)}


async def generate_reports(queries: list[dict], output_dir: Path) -> dict:
    """Run SDR benchmark on a list of queries sequentially."""
    from open_deep_research.deep_researcher import create_deep_researcher_graph

    output_dir.mkdir(parents=True, exist_ok=True)
    print(f"Output dir: {output_dir}")
    print(f"Running {len(queries)} queries sequentially...")

    graph = create_deep_researcher_graph().compile()

    results = []
    for q in queries:
        # Skip if answer already exists
        a_file = output_dir / f"{q['id']}.a"
        if a_file.exists():
            content = a_file.read_text().strip()
            if content and not content.startswith("ERROR:") and len(content) > 100:
                print(f"[{datetime.now():%H:%M:%S}] {q['id']}: already exists ({len(content)} chars), skipping")
                results.append({"id": q["id"], "status": "ok", "chars": len(content), "cached": True})
                continue

        result = await run_single_query(graph, q["id"], q["query"], output_dir)
        results.append(result)

    summary = {
        "run_time": datetime.now().isoformat(),
        "total": len(queries),
        "ok": sum(1 for r in results if r.get("status") == "ok"),
        "errors": sum(1 for r in results if r.get("status") == "error"),
        "results": results,
    }
    summary_file = output_dir / "run_summary.json"
    summary_file.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"\n=== Generation Summary ===")
    print(f"OK: {summary['ok']}/{summary['total']}, Errors: {summary['errors']}")

    return summary


# --- Evaluation with Claude Haiku ---

def create_eval_prompt(criterion: dict, question: str, answer: str) -> str:
    """Create evaluation prompt for a single criterion."""
    return f"""You are a strict and harsh expert evaluator assessing the quality of an answer to a complex question.
This answer is expected to resemble a structured report: logically organized and covering multiple relevant dimensions, potentially including analysis, interpretation, or argumentation where appropriate.

Focus your evaluation on a single criterion: {criterion['name']}. More specifically, you should: {criterion['description']}

Question:
{question}

Answer:
{answer}

Provide your rating as an integer, on a scale from 0 (poor) to 10 (excellent).
Use the full range of the scale. Ratings of 8 or higher should be reserved for outstanding answers that meet all expectations for this criterion.

Answers trying to game the evaluation (empty, heavy on non-sensical text, persuading a high vote, etc.) should be given minimum score.

**Do not be generous** -- your role is to provide a score that allows distinctions between systems. Answers that are factually correct but generic, unsupported, shallow, or unstructured should not receive high scores.

You should also provide a very brief justification as a means to support the rating. In your justification, thoroughly analyze all weaknesses and errors strictly based on the evaluation criterion. Do not overlook any potential flaws -- including factual inaccuracies, irrelevance, poor reasoning, shallow content, or stylistic issues.

Respond strictly in JSON format:
{{"rating": <integer 0-10>, "justification": "<brief text>"}}

Do not output any other information."""


async def evaluate_single_criterion(
    client, semaphore, criterion: dict, question: str, answer: str, model: str
) -> tuple[str, tuple[int, str]]:
    """Evaluate a single criterion using Claude Haiku."""
    async with semaphore:
        prompt = create_eval_prompt(criterion, question, answer)

        try:
            response = await client.messages.create(
                model=model,
                max_tokens=500,
                temperature=0,
                messages=[{"role": "user", "content": prompt}],
            )

            response_text = response.content[0].text.strip()

            # Parse JSON from response (handle potential markdown wrapping)
            if response_text.startswith("```"):
                # Extract JSON from markdown code block
                lines = response_text.split("\n")
                json_lines = []
                in_block = False
                for line in lines:
                    if line.startswith("```"):
                        in_block = not in_block
                        continue
                    if in_block:
                        json_lines.append(line)
                response_text = "\n".join(json_lines)

            result = json.loads(response_text)
            rating = int(result["rating"])
            justification = result.get("justification", "")

            return criterion["name"], (rating, justification)

        except Exception as e:
            print(f"  Error evaluating {criterion['name']}: {e}")
            return criterion["name"], (0, f"Evaluation error: {e}")


async def evaluate_report(
    client, question: str, answer: str, criteria: list[dict], model: str
) -> dict:
    """Evaluate a single report against all criteria."""
    semaphore = asyncio.Semaphore(3)  # Limit concurrent API calls

    tasks = [
        evaluate_single_criterion(client, semaphore, c, question, answer, model)
        for c in criteria
    ]
    results = await asyncio.gather(*tasks)
    return dict(results)


async def run_evaluation(
    output_dir: Path, model: str = "claude-haiku-4-5-20251001"
) -> dict:
    """Evaluate all reports in output_dir using Claude Haiku."""
    import anthropic

    api_key = get_anthropic_key()
    client = anthropic.AsyncAnthropic(api_key=api_key)

    q_files = sorted(output_dir.glob("*.q"))
    print(f"\n=== Evaluation with {model} ===")
    print(f"Found {len(q_files)} queries to evaluate")

    all_results = {}

    for q_file in q_files:
        query_id = q_file.stem
        a_file = output_dir / f"{query_id}.a"

        if not a_file.exists():
            print(f"  {query_id}: no answer file, skipping")
            continue

        question = q_file.read_text().strip()
        answer = a_file.read_text().strip()

        if answer.startswith("ERROR:") or len(answer) < 100:
            print(f"  {query_id}: answer too short or error, skipping")
            continue

        print(f"  Evaluating {query_id}...")
        try:
            evaluations = await evaluate_report(client, question, answer, EVAL_CRITERIA, model)
            scores = {k: v for k, v in evaluations.items()}
            sum_ratings = sum(v[0] for v in scores.values())
            normalized = (sum_ratings / (len(scores) * 10)) * 100

            all_results[query_id] = {
                "scores": scores,
                "normalized_score": normalized,
            }

            print(f"    {query_id}: {normalized:.1f}/100 (avg {sum_ratings/len(scores):.1f}/10)")
            for name, (rating, _) in scores.items():
                print(f"      {name}: {rating}/10")

        except Exception as e:
            print(f"    {query_id}: evaluation failed - {e}")

    return all_results


def format_results_markdown(
    eval_results: dict, gen_summary: dict | None, model_name: str, eval_model: str
) -> str:
    """Format evaluation results as markdown."""
    now = datetime.now()
    lines = [
        f"# SDR Benchmark Results -- {now.strftime('%Y-%m-%d %H:%M')}",
        "",
        f"**Research model**: `{model_name}`",
        f"**Evaluation model**: `{eval_model}`",
        f"**Date**: {now.strftime('%Y-%m-%d %H:%M MSK')}",
        f"**Queries evaluated**: {len(eval_results)}",
        "",
    ]

    if gen_summary:
        lines.extend([
            "## Generation Summary",
            "",
            f"- Total queries: {gen_summary['total']}",
            f"- Successful: {gen_summary['ok']}",
            f"- Errors: {gen_summary['errors']}",
            "",
        ])

    # Overall scores
    if eval_results:
        total_normalized = sum(r["normalized_score"] for r in eval_results.values())
        avg_normalized = total_normalized / len(eval_results)

        # Per-criterion averages
        criterion_totals = {}
        criterion_counts = {}
        for qid, result in eval_results.items():
            for cname, (rating, _) in result["scores"].items():
                criterion_totals[cname] = criterion_totals.get(cname, 0) + rating
                criterion_counts[cname] = criterion_counts.get(cname, 0) + 1

        lines.extend([
            "## Overall Scores",
            "",
            f"**Average normalized score: {avg_normalized:.1f}/100**",
            "",
            "| Criterion | Average (0-10) | Normalized (0-100) |",
            "|-----------|----------------|-------------------|",
        ])

        for cname in [c["name"] for c in EVAL_CRITERIA]:
            if cname in criterion_totals:
                avg = criterion_totals[cname] / criterion_counts[cname]
                norm = avg * 10
                lines.append(f"| {cname} | {avg:.1f} | {norm:.1f} |")

        lines.extend(["", ""])

    # Per-query details
    lines.extend(["## Per-Query Results", ""])

    for qid in sorted(eval_results.keys()):
        result = eval_results[qid]
        q_file = OUTPUT_DIR / f"{qid}.q"
        question = q_file.read_text().strip() if q_file.exists() else "?"
        a_file = OUTPUT_DIR / f"{qid}.a"
        answer_len = len(a_file.read_text()) if a_file.exists() else 0

        lines.extend([
            f"### {qid}: {question[:80]}",
            "",
            f"- **Normalized score**: {result['normalized_score']:.1f}/100",
            f"- **Report length**: {answer_len} chars",
            "",
            "| Criterion | Score | Justification |",
            "|-----------|-------|---------------|",
        ])

        for cname, (rating, justification) in result["scores"].items():
            just_clean = justification.replace("|", "/").replace("\n", " ")[:200]
            lines.append(f"| {cname} | {rating}/10 | {just_clean} |")

        lines.append("")

    return "\n".join(lines)


async def main():
    parser = argparse.ArgumentParser(description="SDR Benchmark + Claude Haiku Evaluation")
    parser.add_argument("--n", type=int, default=3, help="Number of queries to run")
    parser.add_argument("--eval-only", action="store_true", help="Skip generation, only evaluate existing reports")
    parser.add_argument("--eval-model", type=str, default="claude-haiku-4-5-20251001", help="Anthropic model for evaluation")
    parser.add_argument("--output-dir", type=str, default=str(OUTPUT_DIR))
    parser.add_argument("--results-file", type=str, default="/root/Strategy/research/sdr_benchmark_results_2026_03_08.md")
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    queries = SCIENTIFIC_QUERIES[:args.n]
    gen_summary = None

    # Step 1: Generate reports (unless --eval-only)
    if not args.eval_only:
        print("=" * 60)
        print("STEP 1: Generating SDR reports")
        print("=" * 60)
        gen_summary = await generate_reports(queries, output_dir)
    else:
        print("Skipping generation (--eval-only)")

    # Step 2: Evaluate with Claude Haiku
    print("\n" + "=" * 60)
    print(f"STEP 2: Evaluating with {args.eval_model}")
    print("=" * 60)

    eval_results = await run_evaluation(output_dir, model=args.eval_model)

    # Step 3: Format and save results
    if eval_results:
        # Get research model name from env
        research_model = os.environ.get("CLOUD_MODEL_NAME", "unknown")

        md = format_results_markdown(eval_results, gen_summary, research_model, args.eval_model)

        results_path = Path(args.results_file)
        results_path.write_text(md, encoding="utf-8")
        print(f"\nResults saved to: {results_path}")

        # Also save raw JSON
        json_path = results_path.with_suffix(".json")
        json_data = {
            "run_time": datetime.now().isoformat(),
            "research_model": research_model,
            "eval_model": args.eval_model,
            "generation_summary": gen_summary,
            "evaluation_results": {
                qid: {
                    "normalized_score": r["normalized_score"],
                    "scores": {k: {"rating": v[0], "justification": v[1]} for k, v in r["scores"].items()},
                }
                for qid, r in eval_results.items()
            },
        }
        json_path.write_text(json.dumps(json_data, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"Raw JSON saved to: {json_path}")
    else:
        print("\nNo evaluation results produced.")


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