from __future__ import annotations

import os
from functools import lru_cache

from src.envs import PROJECT_ROOT

TASK_README_DIRS = {
    "ruaime": "ruAIME",
    "t_math": "tmath",
    "luzitania": "luzitania",
    "mmred": "mmred",
}


def _readme_search_roots() -> list[str]:
    """Resolve README roots; prefer local MERA_Reason/datasets, then optional overrides."""
    roots = []
    # Local copies shipped with the leaderboard (kept in sync with dataset releases).
    roots.append(os.path.join(PROJECT_ROOT, "datasets"))
    env_root = os.environ.get("MERA_DATASETS_DIR")
    if env_root:
        roots.append(env_root)
    sibling = os.path.join(
        os.path.dirname(PROJECT_ROOT),
        "reasoning",
        "MERA",
        "datasets",
    )
    roots.append(sibling)
    return roots


@lru_cache(maxsize=None)
def _load_readme(task_name: str) -> str:
    folder = TASK_README_DIRS.get(task_name)
    if not folder:
        return f"*Unknown dataset: `{task_name}`*"

    for filename in ("README_ru.md", "README.md"):
        for root in _readme_search_roots():
            path = os.path.join(root, folder, filename)
            if os.path.isfile(path):
                with open(path, encoding="utf-8") as f:
                    return f.read().strip()

    return f"*README not found for `{task_name}`.*"


def get_dataset_readme_markdown(task_name: str) -> str:
    return _load_readme(task_name)