solar_eval/core/project_loader.py

"""Load project configurations from YAML files."""

import logging
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import yaml

from solar_eval.core.pipeline_compose import PipelineCompositionError, compose_pipeline
from solar_eval.core.task_discovery import discover_tasks

logger = logging.getLogger(__name__)


def load_all_project_configs(
    data_dir: str | Path, config_dirs: Mapping[str, Path] | None = None
) -> list[dict[str, Any]]:
    """레포 관리 config 디렉토리 + 데이터 루트를 병합해 로드한다.

    config_dirs (프로젝트명 → accounts/<a>/<c>/03-evaluation) 가 우선이고,
    데이터 루트에서만 발견되는 프로젝트(레포 관리 제외 트랙)는 뒤에 붙는다.
    파이프라인 참조는 각 config 가 놓인 디렉토리 기준으로 해석된다.
    """
    configs: list[dict[str, Any]] = []
    seen: set[str] = set()
    for _, config_dir in sorted((config_dirs or {}).items()):
        config = _load_yaml(Path(config_dir) / "project.yaml")
        if config is None or config.get("name") in seen:
            continue
        _resolve_defaults(config, Path(config_dir), data_root=Path(data_dir))
        configs.append(config)
        seen.add(config["name"])
    for config in load_project_configs(data_dir):
        if config["name"] not in seen:
            configs.append(config)
    return configs


def load_project_configs(projects_dir: str | Path) -> list[dict[str, Any]]:
    """Load all YAML project configs from a directory.

    Searches two locations (both supported, no duplicates):
      1. projects_dir/<name>/project.yaml  (preferred — config inside project folder)
      2. projects_dir/<name>.yaml          (legacy — config at root level)

    Args:
        projects_dir: Path to directory containing project configs.

    Returns:
        List of parsed project config dicts.
    """
    configs = []
    seen_names: set[str] = set()
    projects_path = Path(projects_dir)
    if not projects_path.exists():
        logger.warning(f"Projects directory not found: {projects_path}")
        return configs

    # 1. project.yaml inside project folders (preferred)
    for project_yaml in sorted(projects_path.glob("*/project.yaml")):
        config = _load_yaml(project_yaml)
        if config and config.get("name") not in seen_names:
            project_dir = project_yaml.parent
            _resolve_defaults(config, project_dir, data_root=projects_path)
            configs.append(config)
            seen_names.add(config["name"])

    # 2. Legacy: <name>.yaml at root level
    for yaml_file in sorted(projects_path.glob("*.yaml")):
        config = _load_yaml(yaml_file)
        if config and config.get("name") not in seen_names:
            project_dir = yaml_file.parent / config["name"]
            _resolve_defaults(config, project_dir, data_root=projects_path)
            configs.append(config)
            seen_names.add(config["name"])

    return configs


def _load_yaml(path: Path) -> dict[str, Any] | None:
    """Load and validate a single project YAML file."""
    try:
        with open(path) as f:
            config = yaml.safe_load(f)
        if config and "name" in config:
            logger.debug(f"Loaded project config: {config['name']} from {path}")
            return config
    except Exception as e:
        logger.warning(f"Failed to load {path}: {e}")
    return None


def _resolve_defaults(
    config: dict[str, Any], project_dir: Path, data_root: Path | None = None
) -> None:
    """Merge top-level `defaults` into each task, then resolve derived fields.

    Also expands `discover:` rules into tasks, resolves string `pipeline_config`
    references to actual pipeline YAML files (following `extends`), and attaches a
    resolved `field_map` (see `resolve_field_map`). Mutates config in place.
    """
    # 로더가 채우는 유도 필드 — 스텝이 config 정본 기준 자산(치환 사전 등)을
    # 풀 때 쓴다. project.yaml 에 적는 값이 아니다.
    config["config_dir"] = str(project_dir)

    _expand_discovered_tasks(config, data_root)

    tasks = config.get("tasks", [])
    defaults = config.get("defaults", {})

    for task in tasks:
        # Merge each default key into task if not already set
        for key, value in defaults.items():
            if key not in task:
                task[key] = value

        # Resolve string pipeline_config → load from pipelines/ directory
        _resolve_pipeline_ref(task, project_dir)

        # EvalSample field_map — explicit if declared, else derived from the
        # legacy input_fields/golden_field/golden_fields trio (no-op today: nothing
        # consumes task["field_map"] yet, this only prepares stage C/D wiring).
        task["field_map"] = resolve_field_map(task)


def _expand_discovered_tasks(config: dict[str, Any], data_root: Path | None) -> None:
    """`discover:` 규칙이 찾은 태스크를 `tasks` 뒤에 붙인다 (명시 태스크가 우선).

    발견 실패는 삼키지 않고 로그로 남긴다 — 데이터 루트가 없는 환경(CI 등)에서는
    빈 목록이 정상이지만, 규칙 자체가 잘못된 경우와는 구분돼야 한다.
    """
    if not config.get("discover"):
        return
    try:
        found = discover_tasks(config, data_root)
    except ValueError as e:
        logger.warning("Task discovery failed for %s: %s", config.get("name"), e)
        return
    if found:
        config["tasks"] = [*(config.get("tasks") or []), *found]
        logger.debug("Discovered %d tasks for %s", len(found), config.get("name"))


def load_pipeline_file(pipelines_dir: Path, name: str) -> dict[str, Any]:
    """파이프라인 YAML 한 장을 읽어 `extends`/override 를 해석한 dict 로 돌려준다.

    Args:
        pipelines_dir: `pipelines/` 디렉토리.
        name: 확장자 없는 파이프라인 이름.

    Raises:
        FileNotFoundError: 파일이 없을 때.
        PipelineCompositionError: 상속·override 해석이 실패했을 때.
    """
    path = pipelines_dir / f"{name}.yaml"
    if not path.exists():
        raise FileNotFoundError(f"Pipeline file not found: {path}")
    with open(path) as f:
        raw = yaml.safe_load(f)
    if not isinstance(raw, dict):
        raise PipelineCompositionError(f"Pipeline {name!r} is not a mapping: {path}")
    return compose_pipeline(raw, load_base=lambda base: _load_raw_pipeline(pipelines_dir, base))


def _load_raw_pipeline(pipelines_dir: Path, name: str) -> dict[str, Any]:
    """`extends` 대상을 조립하지 않은 상태로 읽는다 (다단 상속 검출용)."""
    path = pipelines_dir / f"{name}.yaml"
    if not path.exists():
        raise PipelineCompositionError(f"Base pipeline {name!r} not found: {path}")
    with open(path) as f:
        raw = yaml.safe_load(f)
    if not isinstance(raw, dict):
        raise PipelineCompositionError(f"Base pipeline {name!r} is not a mapping: {path}")
    return raw


def _resolve_pipeline_ref(task: dict[str, Any], project_dir: Path) -> None:
    """If pipeline_config is a string reference, load the pipeline YAML file."""
    pipeline_ref = task.get("pipeline_config")
    if not isinstance(pipeline_ref, str):
        return

    try:
        task["pipeline_config"] = load_pipeline_file(project_dir / "pipelines", pipeline_ref)
        logger.debug(f"Resolved pipeline_config '{pipeline_ref}' from {project_dir}")
    except (FileNotFoundError, PipelineCompositionError, OSError, yaml.YAMLError) as e:
        logger.warning(f"Failed to load pipeline '{pipeline_ref}': {e}")


def resolve_field_map(task_config: dict[str, Any]) -> dict[str, Any]:
    """task 설정에서 `EvalSample.from_row` 용 field_map 을 결정한다.

    순수 함수 — task_config 를 변형하지 않고 새 dict 를 반환한다.

    우선순위:
      1. task_config 에 명시적 `field_map` 이 있으면 그대로(사본) 반환한다.
      2. 없으면 레거시 `input_fields`/`golden_field`/`golden_fields` 에서 유도한다:
         - `input_fields` 의 **첫 번째** 필드 → `input`
           (나머지 필드는 field_map 이 담지 못한다 — 다중 입력 필드 태스크는
           명시적 field_map 을 선언해야 한다)
         - `golden_field` → `reference` (문자열 컬럼명)
         - `golden_field` 가 없고 `golden_fields` 가 있으면 → `reference` 에
           그 dict 를 그대로 넣는다 (`EvalSample.from_row` 가 여러 컬럼을 묶어
           합성 정답을 만든다). `runner.py` 의 `golden_fields` 처리와 동치.

    Args:
        task_config: `defaults` 병합이 끝난 task 설정 dict.

    Returns:
        EvalSample 필드명 -> 컬럼명(str) 또는 {서브키: 컬럼명}(dict) 매핑.
        어느 소스도 없으면 빈 dict.
    """
    explicit = task_config.get("field_map")
    if explicit:
        return dict(explicit)

    field_map: dict[str, Any] = {}

    input_fields = task_config.get("input_fields") or []
    if input_fields:
        field_map["input"] = input_fields[0]

    golden_field = task_config.get("golden_field")
    golden_fields = task_config.get("golden_fields")
    if golden_field:
        field_map["reference"] = golden_field
    elif golden_fields:
        field_map["reference"] = dict(golden_fields)

    return field_map