import os
from pathlib import Path
from typing import Any, Union

import yaml

EVAL_KEY = "Eval Listing: Automatically Generated"
OPTIONS_KEY = "Options: Automatically Generated"
USAGE_KEY = "Usage: Automatically Generated"
CONTRIBUTORS_KEY = "Contributors: Automatically Generated"

GROUP_SORT_ORDER = (
    "Coding",
    "Assistants",
    "Cybersecurity",
    "Safeguards",
    "Mathematics",
    "Reasoning",
    "Knowledge",
)


def link_md(text: str, href: str) -> str:
    return f"[{text}]({href})"


def contributor_links(contributors: list[str]) -> list[str]:
    links = [
        link_md(f"@{c.strip()}", f"https://github.com/{c.strip()}")
        for c in contributors
    ]
    return links


def listing_md(listing: dict[str, Any]) -> str:
    # form contributor links
    if "contributors" in listing:
        contributors = f"\n  <sub><sup>Contributed by: {', '.join(contributor_links(listing['contributors']))}</sub></sup>"
    else:
        contributors = ""

    # form the final markdown
    output: list[str] = []
    output.append(
        f"- ### {link_md(listing['title'].strip(), os.path.join(listing['path']))}"
    )
    output.append("")
    output.append(f"  {listing['description'].strip()}{contributors}")
    output.append("")
    output.append("  ```bash")
    for index, task in enumerate(listing["tasks"]):
        if index > 3:
            break
        output.append(f"  inspect eval inspect_evals/{task['name']}")

    output.append("  ```")
    return "\n".join(output)


class Contents:
    def __init__(self, contains_key: bool, prefix: list[str], suffix: list[str]):
        self.contains_key = contains_key
        self.prefix = prefix
        self.suffix = suffix


def readme_contents(file: Path, key: str) -> Contents:
    start_key = f"<!-- {key} -->"
    end_key = f"<!-- /{key} -->"

    # Read the file lines
    readme_lines = []
    with open(file, "r", encoding="utf-8") as readme_file:
        readme_lines = readme_file.readlines()

    # Go through line by line and collect the prefix and suffix
    # to the generated section
    prefix: list[str] = []
    suffix: list[str] = []
    contains_key: bool = False
    collecting: Union[str, None] = "prefix"
    for line in readme_lines:
        line_content = line.rstrip("\r\n")
        if line_content == start_key:
            prefix.append(start_key)
            collecting = None
            contains_key = True
        elif line_content == end_key:
            suffix.append(end_key)
            collecting = "suffix"
        else:
            if collecting == "prefix":
                prefix.append(line_content)
            elif collecting == "suffix":
                suffix.append(line_content)

    return Contents(prefix=prefix, suffix=suffix, contains_key=contains_key)


def rewrite_task_readme(path: str, key: str, contents: list[str]) -> None:
    readme_path = Path(__file__).parent / ".." / path / "README.md"
    parsed = readme_contents(readme_path, key)
    if parsed.contains_key:
        with open(readme_path, "w", encoding="utf-8") as readme_file:
            readme_file.write(
                "\n".join(parsed.prefix + contents + parsed.suffix) + "\n"
            )


def generate_options(task_metadata: dict[str, Any]) -> None:
    task_list = [task["name"] for task in task_metadata["tasks"]]
    task_names = (task_list * 3)[:3]

    contents: list[str] = []
    contents.append("## Options")
    contents.append("")
    contents.append(
        "You can control a variety of options from the command line. For example:"
    )
    contents.append("")
    contents.append("```bash")
    contents.append(f"inspect eval inspect_evals/{task_names[0]} --limit 10")
    contents.append(f"inspect eval inspect_evals/{task_names[1]} --max-connections 10")
    contents.append(f"inspect eval inspect_evals/{task_names[2]} --temperature 0.5")
    contents.append("```")
    contents.append("")
    contents.append("See `inspect eval --help` for all available options.")

    rewrite_task_readme(task_metadata["path"], OPTIONS_KEY, contents)


def generate_usage(task_metadata: dict[str, Any]) -> None:
    dependency = task_metadata["dependency"] if "dependency" in task_metadata else None

    contents: list[str] = []
    contents.append("## Usage")
    contents.append("")
    contents.append(
        "First, install the `inspect_ai` and `inspect_evals` Python packages with:"
    )
    contents.append("")
    contents.append("```bash")
    contents.append("pip install inspect_ai")
    if dependency is None:
        contents.append(
            "pip install git+https://github.com/UKGovernmentBEIS/inspect_evals"
        )
    else:
        contents.append(
            f'pip install "inspect_evals[{dependency}]@git+https://github.com/UKGovernmentBEIS/inspect_evals"'
        )
    contents.append("```")
    contents.append("")

    contents.append(
        "Or, if developing on a clone of the `inspect_evals` repo, you can install the package in editable mode with:"
    )
    contents.append("")
    contents.append("```bash")
    if dependency is None:
        contents.append('pip install -e ".[dev]"')
    else:
        contents.append(f'pip install -e ".[dev,{dependency}]"')
    contents.append("```")
    contents.append("")

    contents.append("Then, evaluate against one or more models with:")
    contents.append("")
    contents.append("```bash")
    for index, task in enumerate(task_metadata["tasks"]):
        if index > 3:
            break
        contents.append(
            f"inspect eval inspect_evals/{task['name']} --model openai/gpt-4o"
        )
    contents.append("```")
    contents.append("")
    contents.append(
        "After running evaluations, you can view their logs using the `inspect view` command:"
    )
    contents.append("")
    contents.append("```bash")
    contents.append("inspect view")
    contents.append("```")
    contents.append("")
    contents.append(
        "If you don't want to specify the `--model` each time you run an evaluation, create a `.env` configuration file in your working directory that defines the `INSPECT_EVAL_MODEL` environment variable along with your API key. For example:"
    )
    contents.append("")
    contents.append("```bash")
    contents.append("INSPECT_EVAL_MODEL=anthropic/claude-3-5-sonnet-20240620")
    contents.append("ANTHROPIC_API_KEY=<anthropic-api-key>")
    contents.append("```")

    rewrite_task_readme(task_metadata["path"], USAGE_KEY, contents)


def generate_contributors(task_metadata: dict[str, Any]) -> None:
    content = []
    if "contributors" in task_metadata:
        content.append(
            f"Contributed by {', '.join(contributor_links(task_metadata['contributors']))}"
        )

    rewrite_task_readme(
        task_metadata["path"],
        CONTRIBUTORS_KEY,
        content,
    )


def generate_readme() -> None:
    # directory configuration
    readme_path = Path(__file__).parent / "../README.md"
    listing_file = (
        Path(__file__).parent.parent / "src" / "inspect_evals" / "listing.yaml"
    )

    # read the listings
    listings_raw: list[dict[str, Any]] = []
    with open(listing_file, "r", encoding="utf-8") as yaml_file:
        listings_raw = yaml.safe_load(yaml_file)

    # read them
    listing_groups: dict[str, list[dict[str, Any]]] = {}
    for listing_raw in listings_raw:
        # make sure required keys are here
        if not all(
            key in listing_raw
            for key in ["title", "description", "path", "group", "tasks"]
        ):
            raise RuntimeError(
                f"This listing.yaml contains this entry which is missing a required field:\n{listing_raw}"
            )

        # place the listing in a group
        if listing_raw["group"] not in listing_groups:
            listing_groups[listing_raw["group"]] = []
        listing_groups[listing_raw["group"]].append(listing_raw)

    # sort the listings within each group by title and path
    for group in listing_groups:
        listing_groups[group] = sorted(
            listing_groups[group], key=lambda x: (x["title"], x["path"])
        )

    # sort the groups by specified order
    # Create a mapping of group name to sort index, with any unlisted groups going to the end
    sort_index = {name: i for i, name in enumerate(GROUP_SORT_ORDER)}
    listing_groups = dict(
        sorted(
            listing_groups.items(),
            key=lambda x: sort_index.get(x[0], len(GROUP_SORT_ORDER)),
        )
    )

    # generate the markdown
    content: list[str] = []
    for group, listings in listing_groups.items():
        content.append(f"## {group}")
        content.append("")
        for listing in listings:
            content.append(listing_md(listing))
            content.append("")

    # write the readme
    contents = readme_contents(readme_path, EVAL_KEY)

    # rewrite the readme with prefix and suffix content
    with open(readme_path, "w", encoding="utf-8") as readme_file:
        readme_file.write("\n".join(contents.prefix + content + contents.suffix) + "\n")

    for listing_raw in listings_raw:
        generate_options(listing_raw)
        generate_usage(listing_raw)
        generate_contributors(listing_raw)


if __name__ == "__main__":
    generate_readme()
