"""
ashr-labs — one command to set up automatic agent testing with Claude Code.

Usage:
    npx ashr-labs <api-key>
    ashr-labs tp_abc123
    ASHR_LABS_API_KEY=tp_... ashr-labs
"""

from __future__ import annotations

import json
import os
import re
import sys
from pathlib import Path

# ANSI helpers
ASHR_BLUE = "\033[38;5;69m"
DIM = "\033[2m"
BOLD = "\033[1m"
RESET = "\033[0m"
GREEN = "\033[32m"
YELLOW = "\033[33m"


# ──────────────────────────────────────────────────────────────
# Config
# ──────────────────────────────────────────────────────────────

class ProjectConfig:
    __slots__ = (
        "api_key_env_var", "agent_name", "agent_description",
        "entrypoint", "domain", "lang",
    )

    def __init__(
        self, *, api_key_env_var: str, agent_name: str,
        agent_description: str, entrypoint: str, domain: str, lang: str,
    ) -> None:
        self.api_key_env_var = api_key_env_var
        self.agent_name = agent_name
        self.agent_description = agent_description
        self.entrypoint = entrypoint
        self.domain = domain
        self.lang = lang


# ──────────────────────────────────────────────────────────────
# File generators
# ──────────────────────────────────────────────────────────────

def generate_ashr_config(cfg: ProjectConfig) -> str:
    return json.dumps(
        {
            "$schema": "https://rank.ashr.io/schemas/ashr.json",
            "apiKeyEnvVar": cfg.api_key_env_var,
            "agent": {
                "name": cfg.agent_name,
                "description": cfg.agent_description,
                "entrypoint": cfg.entrypoint,
                "domain": cfg.domain,
            },
            "eval": {"autoRun": True, "scenarioCount": 5, "maxWorkers": 3},
        },
        indent=2,
    )


def generate_claude_md_section(cfg: ProjectConfig) -> str:
    is_ts = cfg.lang == "typescript"
    sdk_import = (
        'import { AshrLabsClient, EvalRunner } from "ashr-labs";'
        if is_ts
        else "from ashr_labs import AshrLabsClient, EvalRunner"
    )
    client_ctor = (
        f"new AshrLabsClient(process.env.{cfg.api_key_env_var}!)"
        if is_ts
        else f'AshrLabsClient(os.environ["{cfg.api_key_env_var}"])'
    )
    lang_tag = "typescript" if is_ts else "python"

    return f"""
# Ashr Labs — Automated Agent Testing

This project uses [Ashr Labs](https://rank.ashr.io) for continuous agent evaluation.
Configuration is in `.ashr.json`. The API key is in the `{cfg.api_key_env_var}` environment variable.

## When to test

Run `/test-agent` or proactively trigger a test whenever you:
- Change agent logic, prompts, or system instructions
- Add, modify, or remove tools the agent can call
- Change how the agent handles edge cases
- Before committing agent-related changes

## How to test

**Prefer reusing an existing dataset** — only generate a new one if the agent's
tools, inputs, or domain have changed. Reusing datasets lets you compare runs
against the same scenarios over time.

```{lang_tag}
{sdk_import}

client = {client_ctor}

# Option 1: Reuse an existing dataset (preferred)
runner = EvalRunner.from_dataset(client, dataset_id=DATASET_ID)

# Option 2: Generate a new dataset (only when agent config changes)
# dataset_id, source = client.generate_dataset(
#     "Eval: {cfg.agent_name}",
#     {{
#         "agent": {{"name": "{cfg.agent_name}", "description": "{cfg.agent_description}"}},
#         "context": {{"domain": "{cfg.domain}"}},
#     }},
# )
# runner = EvalRunner(source)

results = runner.run(agent)
results.deploy(client, dataset_id=DATASET_ID)
```

## Interpreting results

All scoring is performed server-side after `deploy()`. Check the dashboard for:
- `tests_passed / total_tests` — overall pass rate
- `average_similarity_score` — how close agent responses match expected (0-1)
- `total_tool_call_divergence` — count of wrong/missing tool calls
- If pass rate drops after a change, investigate before committing
"""


def generate_test_agent_command(cfg: ProjectConfig) -> str:
    is_ts = cfg.lang == "typescript"
    sdk_pkg = "ashr-labs" if is_ts else "ashr_labs"

    return f"""---
description: Run the Ashr Labs eval suite against the agent and report results.
---

Run an automated evaluation of the "{cfg.agent_name}" agent using the Ashr Labs SDK.
Do this FULLY AUTONOMOUSLY — do not ask the user for input at any step.

## Dataset strategy

**Reuse existing datasets whenever possible.** Only generate a new dataset when:
- The agent's tools have changed (added, removed, or modified)
- The agent's domain or accepted inputs have changed
- You explicitly need fresh/different test scenarios

To find existing datasets, call `client.list_datasets()` and pick the most recent
one for this agent. Store the dataset ID in `.ashr.json` under `eval.datasetId`
for future runs.

## Steps

1. Read `.ashr.json` for project configuration.
2. Read the agent code at `{cfg.entrypoint}` to understand current behavior.
3. Check if `.ashr.json` has `eval.datasetId` — if so, reuse that dataset.
   If not, call `client.list_datasets()` to find an existing one for this agent.
   Only call `client.generate_dataset()` if no suitable dataset exists.
4. Write a SINGLE self-contained eval script (`_ashr_eval.py`) that does EVERYTHING
   in one process — no separate scripts, no separate polling commands:
   - Imports `AshrLabsClient`, `EvalRunner` from `{sdk_pkg}`
   - Imports `time`, `json`, `os`
   - Creates a client using `os.environ["{cfg.api_key_env_var}"]`
   - Uses `EvalRunner.from_dataset(client, dataset_id=ID)` to load the dataset
   - Implements a lightweight Agent wrapper around the actual agent code in `{cfg.entrypoint}`
   - Runs `runner.run(agent)` with `max_workers=1` (sequential — avoids deepcopy issues)
   - Deploys results with `run.deploy(client, dataset_id=ID)`
   - **Polls for grading inline**: loops calling `client.get_run(run_id)` every 20 seconds
     until `aggregate_metrics.tests_passed` is not None (max 5 minutes)
   - Prints final graded results: passed, failed, per-test status, tool mismatches
5. Run the eval script with a long timeout (e.g. 10 minutes).
6. If a new dataset was generated, update `.ashr.json` with the new `eval.datasetId`.
7. Delete `_ashr_eval.py`.

## CRITICAL: Single command execution

The eval script MUST be a single file that runs as a single command. Do NOT write
separate scripts for polling or analysis. Do NOT use `sleep` in bash. Everything
happens inside the Python script so it only requires ONE permission grant to run.

## Important

- Use the `{cfg.api_key_env_var}` env var for the API key — never hardcode it.
- All grading is done server-side — do not compute scores locally.
- Do NOT ask the user for permission or input. Run everything automatically.
- Set a bash timeout of at least 600000ms (10 minutes) when running the eval script.
"""


def generate_improve_agent_command(cfg: ProjectConfig) -> str:
    is_ts = cfg.lang == "typescript"
    sdk_pkg = "ashr-labs" if is_ts else "ashr_labs"

    return f"""---
description: Automatically run evals, analyze failures, fix the agent, and re-test until passing.
---

Continuously improve "{cfg.agent_name}" by running evaluations, analyzing failures, and
applying fixes — ALL AUTOMATICALLY without asking for user input.

## Overview

This is an autonomous improvement loop. You will:
1. Run eval + wait for grading (single script, single command)
2. Analyze failures from the output
3. Fix the agent code
4. Re-run eval (same single script, single command)
5. Repeat until target pass rate is met

Do NOT ask the user for permission between iterations. Just run the loop.

## CRITICAL: Single-script execution

ALL eval + grading + analysis MUST happen in ONE Python script (`_ashr_eval.py`) run as
ONE bash command. This avoids repeated permission prompts. The script must:

1. Import the agent from `{cfg.entrypoint}` and `AshrLabsClient`, `EvalRunner` from `{sdk_pkg}`
2. Load dataset from `.ashr.json` `eval.datasetId` using `EvalRunner.from_dataset()`
3. Run `runner.run(agent)` with `max_workers=1`
4. Deploy with `run.deploy(client, dataset_id=ID)`
5. Poll `client.get_run(run_id)` every 20s until `aggregate_metrics.tests_passed` is not None (max 5min)
6. Print graded results: pass/fail per test, tool mismatches with expected vs actual args
7. For each FAILED test, also fetch and print the dataset scenario actions so the failure
   context is visible (user messages, expected tool calls at each action index)

Run this script with `timeout: 600000` (10 minutes). The output gives you everything
needed to diagnose failures without running additional commands.

## Iteration loop

### After each eval run, read the output and analyze failures:

Common failure patterns:
- **Tool called at wrong step**: Agent calls the right tool but too early/late.
  Fix: adjust system prompt to be more/less eager about that tool.
- **Tool not called**: Agent didn't call an expected tool.
  Fix: strengthen prompt guidance about when to use that tool.
- **Extra unexpected tool call**: Agent called a tool it shouldn't have.
  Fix: add prompt constraints about when NOT to call that tool.
- **Wrong tool arguments**: Agent called the right tool with wrong args.
  Fix: improve tool descriptions or add examples in the system prompt.
- **Tool retry timing**: Agent auto-retries a failed tool instead of returning control
  (or vice versa). Fix: adjust the tool execution loop in the agent code.

### Apply fixes to `{cfg.entrypoint}`:
- **System prompt changes** fix most failures — adjust instructions about tool timing
- **Tool loop logic** changes fix retry/timing issues
- Make the SMALLEST change that addresses each failure. Do not refactor unrelated code.

### Re-run eval:
After editing `{cfg.entrypoint}`, run the SAME `_ashr_eval.py` script again (it imports
the agent fresh each time). Same single command, same timeout.

### Stop conditions:
- All tests pass or pass rate >= 80%: Stop. Print summary of changes + before/after metrics.
- 5 iterations reached: Stop. Summarize what's still failing.
- No improvement after 2 consecutive iterations: Stop. Explain blockers.

Clean up `_ashr_eval.py` when done.

## Rules

- Do NOT ask the user for input at any step. Run everything automatically.
- Do NOT write separate scripts for polling or analysis — everything in ONE script.
- Do NOT use `sleep` in bash — all waiting happens inside the Python script.
- Set bash timeout to 600000ms (10 min) for each eval run.
- Use `{cfg.api_key_env_var}` env var for the API key.
- All grading is server-side. Never grade or score locally.
"""


def generate_run_eval_script(cfg: ProjectConfig) -> str:
    return f'''"""
Auto-generated eval script for {cfg.agent_name}.
Runs eval, deploys, polls for grading, and prints results — all in one shot.

Usage:
    python run_eval.py
"""

import json
import os
import sys
import time

# --- EDIT THIS SECTION ---
# Import your agent class here. It must have respond(message) -> dict and reset() -> None.
# respond() returns {{"text": "...", "tool_calls": [{{"name": "...", "arguments": {{...}}}}]}}
from agent import SupportAgent as MyAgent  # <-- Replace with your agent class

DATASET_ID = None  # Set after first run, or put in .ashr.json under eval.datasetId
# --- END EDIT ---


def load_config():
    """Load dataset ID from .ashr.json if available."""
    global DATASET_ID
    if DATASET_ID is not None:
        return
    try:
        with open(".ashr.json") as f:
            config = json.load(f)
        DATASET_ID = config.get("eval", {{}}).get("datasetId")
    except (FileNotFoundError, json.JSONDecodeError):
        pass


def main():
    from ashr_labs import AshrLabsClient, EvalRunner

    api_key = os.environ.get("{cfg.api_key_env_var}")
    if not api_key:
        print("ERROR: {cfg.api_key_env_var} not set")
        sys.exit(1)

    client = AshrLabsClient(api_key=api_key)
    load_config()

    if DATASET_ID is None:
        print("ERROR: No dataset ID. Set DATASET_ID in this script or eval.datasetId in .ashr.json")
        sys.exit(1)

    print(f"Loading dataset {{DATASET_ID}}...")
    runner = EvalRunner.from_dataset(client, dataset_id=DATASET_ID)
    print("Dataset loaded.\\n")

    agent = MyAgent()
    run = runner.run(
        agent,
        max_workers=1,
        on_scenario=lambda sid, s: print(f"  Running: {{s.get(\'title\', sid)}}"),
    )

    print("\\nDeploying results...")
    created = run.deploy(client, dataset_id=DATASET_ID)
    run_id = created["id"]
    print(f"Run #{{run_id}} submitted.")

    # Poll for grading
    print(f"\\nWaiting for grading...")
    graded = client.poll_run(
        run_id,
        on_poll=lambda elapsed, r: print(f"  ...grading in progress ({{elapsed}}s)"),
    )

    # Print results
    result = graded.get("result", {{}})
    metrics = result.get("aggregate_metrics", {{}})
    tests = result.get("tests", [])

    passed = metrics.get("tests_passed", "?")
    failed = metrics.get("tests_failed", "?")
    tool_div = metrics.get("total_tool_call_divergence", 0)
    resp_div = metrics.get("total_response_divergence", 0)

    print(f"\\n{{"=" * 60}}")
    print(f"  RESULTS: {{passed}} passed / {{failed}} failed ({{len(tests)}} total)")
    print(f"  Tool divergences: {{tool_div}} | Response divergences: {{resp_div}}")
    print(f"{{"=" * 60}}")

    for t in tests:
        status = t.get("status", "?")
        tid = t.get("test_id", "?")
        marker = "PASS" if status == "completed" else "FAIL"
        print(f"  [{{marker}}] {{tid}}")
        if status == "failed":
            for ar in t.get("action_results", []):
                if ar.get("action_type") == "tool_call":
                    for tc in ar.get("tool_calls", []):
                        ms = tc.get("match_status", "")
                        if ms in ("mismatch", "partial"):
                            exp = tc.get("expected", {{}})
                            act = tc.get("actual", {{}})
                            print(f"    {{ms}}: expected={{exp.get(\'name\',\'?\')}}({{json.dumps(exp.get(\'arguments\', {{}}))}}) "
                                  f"actual={{act.get(\'name\',\'?\')}}({{json.dumps(act.get(\'arguments\', {{}}))}})")
                elif ar.get("action_type") == "text" and ar.get("actor") == "agent":
                    ms = ar.get("match_status", "")
                    if ms in ("mismatch", "missed"):
                        sim = ar.get("semantic_similarity", 0)
                        print(f"    text {{ms}} (similarity: {{sim}})")
    print()


if __name__ == "__main__":
    main()
'''


def generate_hook_settings(cfg: ProjectConfig) -> dict:
    return {
        "hooks": {
            "Stop": [
                {
                    "matcher": "",
                    "hooks": [
                        {
                            "type": "prompt",
                            "prompt": (
                                f'If you modified agent code in this conversation '
                                f'(files related to "{cfg.agent_name}" or {cfg.entrypoint}), '
                                f'remind the user they can run /test-agent to verify the changes. '
                                f'Keep it to one sentence.'
                            ),
                        },
                    ],
                },
            ],
        },
    }


# ──────────────────────────────────────────────────────────────
# File writers
# ──────────────────────────────────────────────────────────────

def _print(msg: str) -> None:
    print(msg)


def write_file(file_path: str, content: str) -> bool:
    p = Path(file_path)
    p.parent.mkdir(parents=True, exist_ok=True)
    if p.exists():
        _print(f"  {YELLOW}skip{RESET}  {file_path}")
        return False
    p.write_text(content, encoding="utf-8")
    _print(f"  {GREEN}create{RESET}  {file_path}")
    return True


def append_to_file(file_path: str, content: str, marker: str) -> bool:
    p = Path(file_path)
    if p.exists():
        existing = p.read_text(encoding="utf-8")
        if marker in existing:
            _print(f"  {YELLOW}skip{RESET}  {file_path}")
            return False
        with p.open("a", encoding="utf-8") as f:
            f.write("\n" + content)
        _print(f"  {GREEN}append{RESET}  {file_path}")
    else:
        p.write_text(content, encoding="utf-8")
        _print(f"  {GREEN}create{RESET}  {file_path}")
    return True


def merge_json_file(file_path: str, new_settings: dict) -> bool:
    p = Path(file_path)
    p.parent.mkdir(parents=True, exist_ok=True)
    existing: dict = {}
    if p.exists():
        try:
            existing = json.loads(p.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            existing = {}
    merged = {**existing}
    if "hooks" in new_settings:
        existing_hooks = existing.get("hooks", {})
        new_hooks = new_settings["hooks"]
        merged["hooks"] = {**existing_hooks}
        for event, handlers in new_hooks.items():
            merged["hooks"][event] = existing_hooks.get(event, []) + handlers
    p.write_text(json.dumps(merged, indent=2) + "\n", encoding="utf-8")
    _print(f"  {GREEN}create{RESET}  {file_path}")
    return True


# ──────────────────────────────────────────────────────────────
# Auto-detection
# ──────────────────────────────────────────────────────────────

def detect_lang() -> str:
    if os.path.exists("pyproject.toml") or os.path.exists("setup.py") or os.path.exists("requirements.txt"):
        return "python"
    if os.path.exists("tsconfig.json") or os.path.exists("package.json"):
        return "typescript"
    return "python"


def detect_entrypoint(lang: str) -> str:
    candidates = (
        ["agent.py", "src/agent.py", "app/agent.py", "main.py", "src/main.py"]
        if lang == "python"
        else ["src/agent.ts", "src/index.ts", "agent.ts", "index.ts"]
    )
    for c in candidates:
        if os.path.exists(c):
            return c
    return "agent.py" if lang == "python" else "src/agent.ts"


def detect_agent_name() -> str:
    # Try package.json
    if os.path.exists("package.json"):
        try:
            with open("package.json") as f:
                pkg = json.load(f)
            if pkg.get("name"):
                return pkg["name"]
        except (json.JSONDecodeError, OSError):
            pass
    # Try pyproject.toml
    if os.path.exists("pyproject.toml"):
        try:
            with open("pyproject.toml") as f:
                toml = f.read()
            m = re.search(r'^name\s*=\s*"([^"]+)"', toml, re.MULTILINE)
            if m:
                return m.group(1)
        except OSError:
            pass
    # Directory name
    return Path.cwd().name


def detect_description() -> str:
    if os.path.exists("package.json"):
        try:
            with open("package.json") as f:
                pkg = json.load(f)
            if pkg.get("description"):
                return pkg["description"]
        except (json.JSONDecodeError, OSError):
            pass
    if os.path.exists("pyproject.toml"):
        try:
            with open("pyproject.toml") as f:
                toml = f.read()
            m = re.search(r'^description\s*=\s*"([^"]+)"', toml, re.MULTILINE)
            if m:
                return m.group(1)
        except OSError:
            pass
    return "An AI agent"


# ──────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────

def main() -> None:
    args = sys.argv[1:]

    # Help
    if "--help" in args or "-h" in args:
        _print(f"\n{BOLD}ashr-labs{RESET} — set up automatic agent testing\n")
        _print(f"Usage: ashr-labs <api-key>\n")
        _print(f"  The API key is the only required argument. Everything else")
        _print(f"  is auto-detected from your project.\n")
        _print(f"  You can also set ASHR_LABS_API_KEY in your environment")
        _print(f"  and run {DIM}ashr-labs{RESET} with no arguments.\n")
        _print(f"  Get your key at {BOLD}https://app.ashr.io → API Keys{RESET}\n")
        sys.exit(0)

    # Find API key: first arg that starts with tp_, or env var
    # Skip "init" for backward compat
    positional = [a for a in args if a != "init" and not a.startswith("-")]
    api_key = next((a for a in positional if a.startswith("tp_")), "") or os.environ.get("ASHR_LABS_API_KEY", "")

    if not api_key:
        _print(f"\n{BOLD}{ASHR_BLUE}  ashr labs{RESET}\n")
        _print(f"  Usage: {BOLD}ashr-labs <api-key>{RESET}\n")
        _print(f"  Get your key at https://app.ashr.io → API Keys\n")
        sys.exit(1)

    if not api_key.startswith("tp_"):
        _print(f"\n{YELLOW}  Invalid API key — must start with tp_{RESET}")
        _print(f"  Get one at https://app.ashr.io → API Keys\n")
        sys.exit(1)

    _print("")
    _print(f"{BOLD}{ASHR_BLUE}  ashr labs{RESET} {DIM}— setting up agent testing{RESET}")
    _print("")

    # Already initialized?
    if os.path.exists(".ashr.json"):
        _print(f"  {YELLOW}.ashr.json already exists{RESET} — delete it first to re-initialize.\n")
        sys.exit(0)

    # Validate
    _print(f"  {DIM}Validating...{RESET}")
    try:
        from .client import AshrLabsClient
        client = AshrLabsClient(api_key=api_key)
        session = client.init()
        tenant = session.get("tenant", {})
        tenant_name = tenant.get("tenant_name", tenant.get("name", "unknown"))
        _print(f"  {GREEN}✓{RESET} {tenant_name}\n")
    except Exception as e:
        _print(f"  {YELLOW}✗ {e}{RESET}")
        _print(f"  {DIM}Continuing anyway.{RESET}\n")

    # Auto-detect everything
    lang = detect_lang()
    cfg = ProjectConfig(
        api_key_env_var="ASHR_LABS_API_KEY",
        agent_name=detect_agent_name(),
        agent_description=detect_description(),
        entrypoint=detect_entrypoint(lang),
        domain="general",
        lang=lang,
    )

    # Write all files
    write_file(".ashr.json", generate_ashr_config(cfg))

    # .env
    env_path = Path(".env")
    env_line = f"{cfg.api_key_env_var}={api_key}\n"
    if env_path.exists():
        env_content = env_path.read_text(encoding="utf-8")
        if cfg.api_key_env_var not in env_content:
            with env_path.open("a", encoding="utf-8") as f:
                f.write(env_line)
            _print(f"  {GREEN}append{RESET}  .env")
        else:
            _print(f"  {YELLOW}skip{RESET}  .env")
    else:
        env_path.write_text(env_line, encoding="utf-8")
        _print(f"  {GREEN}create{RESET}  .env")

    # .gitignore
    gitignore = Path(".gitignore")
    if gitignore.exists():
        gi = gitignore.read_text(encoding="utf-8")
        if ".env" not in gi:
            with gitignore.open("a", encoding="utf-8") as f:
                f.write("\n.env\n")

    write_file("run_eval.py", generate_run_eval_script(cfg))
    append_to_file("CLAUDE.md", generate_claude_md_section(cfg), "# Ashr Labs")
    write_file(".claude/commands/test-agent.md", generate_test_agent_command(cfg))
    write_file(".claude/commands/improve-agent.md", generate_improve_agent_command(cfg))
    merge_json_file(".claude/settings.json", generate_hook_settings(cfg))

    # Done
    _print(f"\n{GREEN}  Done.{RESET} Open Claude Code and type {BOLD}/test-agent{RESET}\n")


if __name__ == "__main__":
    main()
