"""`extension_grouping` task generator."""
import json
import random
import shutil
from collections import defaultdict
from pathlib import Path

from ..base import (
    Generator, _VERIFY_HEADER, _render_verify, pad_to, _write, _para,
    _AUTHORS, _SONG_TITLES, _STUDENTS, _WORDS,
)


class ExtensionGrouping(Generator):
    KEY = "extension_grouping"
    CATEGORY_NAME = "Extension Grouping"
    DIFFICULTY = "L2"
    TAGS = ["file organization", "file property"]

    EXTS = [".txt", ".md", ".csv", ".json", ".log"]

    def _folder(self, name: str) -> str:
        return Path(name).suffix.lower().lstrip(".") + "_files"

    def build(self, env_dir, llm, rng):
        snippets = llm.gen_snippets("a software project's working files", 9)
        files = []
        for i, snip in enumerate(snippets):
            ext = self.EXTS[i % len(self.EXTS)]
            name = snip["filename"] + ext
            body = snip["content"]
            if ext == ".csv":
                body = "name,value\n" + body.replace(" ", ",")[:60]
            elif ext == ".json":
                body = json.dumps({"note": body[:60]})
            _write(env_dir / name, body)
            files.append(name)
        return {"files": files}

    def description(self, spec):
        return (
            "Please use FileSystem tools to finish the following task:\n\n"
            "### Task: Group files by extension\n\n"
            "For every distinct file extension present in the test directory, create "
            "a folder named `<ext>_files` (the extension lowercased, without the dot — "
            "for example `.txt` files go into `txt_files/`, `.md` into `md_files/`).\n\n"
            "- Move every file into the folder matching its extension.\n"
            "- After you are done, the test directory root must contain only those "
            "folders and no loose files.\n"
            "- Do not rename or modify any file."
        )

    def verify_src(self, spec):
        consts = {"FILES": spec["files"]}
        return _VERIFY_HEADER + f'''

C = json.loads({json.dumps(json.dumps(consts))})


def folder_for(name):
    suffix = name.rsplit(".", 1)[-1].lower() if "." in name else ""
    return suffix + "_files"


def main():
    t = get_test_dir()

    loose = [f.name for f in t.iterdir() if f.is_file() and f.name not in SYSTEM_FILES]
    if loose:
        fail(f"files still in root: {{loose}}")
    ok("root has no loose files")

    located = {{}}
    for d in t.iterdir():
        if d.is_dir():
            for f in d.iterdir():
                if f.is_file() and f.name not in SYSTEM_FILES:
                    located.setdefault(f.name, []).append(d.name)

    for name in C["FILES"]:
        where = located.get(name)
        if not where:
            fail(f"file missing after grouping: {{name}}")
        if len(where) > 1:
            fail(f"file duplicated across folders: {{name}} -> {{where}}")
        expected = folder_for(name)
        if where[0] != expected:
            fail(f"{{name}} is in {{where[0]}}, expected {{expected}}")
    ok(f"all {{len(C['FILES'])}} files grouped correctly by extension")
    print("\\U0001f389 All checks passed!")
    sys.exit(0)


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

    def solve(self, work_dir, spec):
        for name in spec["files"]:
            folder = self._folder(name)
            (work_dir / folder).mkdir(exist_ok=True)
            shutil.move(str(work_dir / name), str(work_dir / folder / name))