import os
from pathlib import Path
import re
import shutil
from typing import Any

import yaml

def generate_site_files() -> None:
    # copy contributor guide
    with open(Path(__file__).parent.parent.parent / "CONTRIBUTING.md", "rt") as f:
        contributing = f.read()
        contributing = contributing.replace("# Contribution Guide", "---\ntitle: Contribution Guide\n---", 1)

    contributing_dir = Path(__file__).parent.parent / "contributing"
    os.makedirs(contributing_dir, exist_ok=True)
    with open(contributing_dir / "index.qmd", "wt") as f:
        f.write(contributing)
    

    # directory configuration
    listing_file = Path(__file__).parent / "../../src/inspect_evals/listing.yaml"

    # read the listings
    listing_groups = read_listings(listing_file)

    # Go through and create directories
    sort_index = 1
    site_dir = Path(__file__).parent.parent
    for group, listings in listing_groups.items():

    # create a directory for the contents
        group_dir = site_dir / "evals" / group.lower()
        if os.path.exists(group_dir):
            shutil.rmtree(group_dir)
        os.makedirs(group_dir)

        # for each listing, create a subdir and place the readme
        for listing in listings:
            # The target directories
            listing_to = group_dir / os.path.basename(listing["path"])
            listing_from = Path(__file__).parent / ".." / ".." / listing["path"].lower()
            
            listing_readme_in = listing_from / "README.md"
            listing_readme_out = listing_to / "index.qmd"

            # Start the final output with fresh front matter
            front_matter = create_front_matter(listing, sort_index)
            sort_index = sort_index + 1

            # Read the README input
            readme_lines = []
            with open(listing_readme_in, "r") as readme_file:
                readme_lines = readme_file.readlines()

            # process the README
            readme_content = process_readme(readme_lines)

            # write the output file
            os.mkdir(listing_to)
            with open(listing_readme_out, "w") as readme_file:
                readme_file.write("\n".join(front_matter + readme_content))

            # snag any other images
            copy_images(listing_from, listing_to)


def read_listings(listing_file: Path) -> dict[str, list[dict[str, Any]]]:

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

    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)

    return listing_groups

def copy_images(listing_from: str, listing_to: str) -> None:
    # snag any other images
    for root, _dirs, files in os.walk(listing_from):
        # Convert root to a Path object
        root_path = Path(root)
        
        # Collect only image files in the current directory
        image_files = [file for file in files if file.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff'))]

        # If there are image files, create the corresponding subdirectory in the target
        if image_files:
            relative_path = root_path.relative_to(listing_from)
            target_subdir = Path(listing_to) / relative_path

            # Create the target subdirectory if it doesn't exist
            target_subdir.mkdir(parents=True, exist_ok=True)

            # Copy each image file to the target subdirectory
            for file in image_files:
                source_file = root_path / file
                target_file = target_subdir / file
                shutil.copy2(source_file, target_file)

def create_front_matter(listing: dict[str, Any], sort_index: int) -> list[str]:

    categories = [listing['group']]
    if 'tags' in listing:
        categories = categories + listing['tags']
    authors = listing['contributors'] if 'contributors' in listing else None

    # fixup the README
    readme_out: list[str] = []
    readme_out.append("---")
    readme_out.append(f"title: \"{listing['title'].strip()}\"")
    readme_out.append(f"description: |\n  {listing['description'].strip()}")
    readme_out.append(f"format: html")
    readme_out.append(f"categories:")
    for category in categories:
        readme_out.append(f"  - \"{category}\"")
    readme_out.append(f"author-title: Contributed By")
    if authors is not None:
        readme_out.append("author:")
        for author in authors:
            readme_out.append(f"  - \"{to_author_link(author)}\"")
    readme_out.append(f"code: {listing['path']}")
    readme_out.append(f"code-url: https://github.com/UKGovernmentBEIS/inspect_evals/tree/main/{listing['path']}")
    if "arxiv" in listing:
        readme_out.append(f"arxiv: {listing['arxiv']}")
    readme_out.append(f"group: {listing['group']}")
    readme_out.append(f"order: {sort_index}")
    readme_out.append("---")
    return readme_out

def to_author_link(author: str) -> str:
    return f"[@{author}](https://github.com/{author})"

def process_readme(readme_lines: list[str]) -> list[str]:
    readme_out: list[str] = []
    eat_title = True
    eating_contributors = False
    for readme_line in readme_lines:
        # Clean the title
        if eat_title and re.match(r'^#\s+', readme_line):
            readme_out.append("## Overview")
            eat_title = False
            continue

        # Clean contributors out altogether
        if not eating_contributors:
            if readme_line.strip().startswith("<!-- Contributors:"):
                eating_contributors = True
                continue
        else:
            eating_contributors = not readme_line.strip().startswith("<!-- /Contributors:")
            continue

        # append the line
        readme_out.append(readme_line.rstrip("\r\n"))
    return readme_out


if __name__ == "__main__":
    generate_site_files()
