"""Task checker — validates and optionally fixes task directories.

Called by: cli.py (check command), also runnable standalone

Checks:
  0. Required files present (task.toml, instruction.md, solution/solve.sh,
     data/needles.py, environment/Dockerfile, tests/evaluate.py)
  1. tests/evaluate.py: importable, has callable evaluate(args...)
  2. tests/evaluate.py: returns "score" or "reward" key (static check)
  3. Dockerfile conventions: base image, log dirs, no hardcoded env vars
  4. No inline skills (must use top-level skills/ dir)
  5. No legacy docker-compose (base image runs envs in-process)
  6. task.toml schema: no unknown fields
  7. config.toml cross-check: Dockerfile.base deps & env vars

Usage:
    clawbench check --tasks-dir tasks/
    clawbench check --task send-data-report
"""

from __future__ import annotations

import argparse
import importlib.util
import inspect
import re
import sys
from pathlib import Path

try:
    import tomllib
except ModuleNotFoundError:  # Python < 3.11
    import tomli as tomllib  # type: ignore[no-redef]

REPO_ROOT = Path(__file__).parents[3]
TASKS_DIR = REPO_ROOT / "tasks"
SKIP = {"ambiguous-cleanup", "targeted-promo-delete", "vendor-report-organize"}

BASE_IMAGE_PATTERN = re.compile(r"^FROM\s+.*smolclaws-base", re.MULTILINE)


# ---------------------------------------------------------------------------
# Base-image helpers (env-in-base detection)
# ---------------------------------------------------------------------------

def _envs_in_base_image() -> set[str]:
    """Return set of env package names installed in Dockerfile.base."""
    dockerfile_base = REPO_ROOT / "docker" / "Dockerfile.base"
    if not dockerfile_base.exists():
        return set()
    src = dockerfile_base.read_text()
    return set(re.findall(r"packages/environments/(claw-\w+)", src))


def _dockerfile_env_packages(dockerfile_src: str) -> set[str]:
    """Return env package names COPY'd in a task Dockerfile."""
    return set(re.findall(r"packages/environments/(claw-\w+)", dockerfile_src))


# ---------------------------------------------------------------------------
# Check helpers
# ---------------------------------------------------------------------------

HARBOR_REQUIRED_FILES = [
    "task.toml",
    "instruction.md",
    "solution/solve.sh",
    "data/needles.py",
    "environment/Dockerfile",
    "tests/evaluate.py",
]


def check_format(task_dir: Path) -> list[str]:
    """Check that all required task files are present."""
    issues = []
    for rel_path in HARBOR_REQUIRED_FILES:
        if not (task_dir / rel_path).exists():
            issues.append(f"task format: missing {rel_path}")
    return issues


def check_evaluate_callable(task_dir: Path) -> list[str]:
    issues = []
    evaluate_path = task_dir / "tests" / "evaluate.py"
    if not evaluate_path.exists():
        return ["tests/evaluate.py: missing"]

    module_name = f"_check_{task_dir.name}"
    spec = importlib.util.spec_from_file_location(module_name, evaluate_path)
    try:
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
    except Exception as exc:
        return [f"tests/evaluate.py: import failed — {exc}"]
    finally:
        sys.modules.pop(module_name, None)

    fn = getattr(module, "evaluate", None)
    if fn is None or not callable(fn):
        return ["tests/evaluate.py: no callable evaluate()"]

    try:
        params = list(inspect.signature(fn).parameters)
        if not params:
            issues.append("tests/evaluate.py: evaluate() takes no arguments")
    except (ValueError, TypeError):
        pass

    return issues


def check_rewards_contract(task_dir: Path) -> list[str]:
    evaluate_path = task_dir / "tests" / "evaluate.py"
    if not evaluate_path.exists():
        return []
    src = evaluate_path.read_text()
    has_key = '"score"' in src or "'score'" in src or '"reward"' in src or "'reward'" in src
    if not has_key:
        return ["tests/evaluate.py: no 'score' or 'reward' key found (rewards contract)"]
    return []


# ---------------------------------------------------------------------------
# Dockerfile convention checks
# ---------------------------------------------------------------------------

def check_dockerfile_conventions(task_dir: Path) -> tuple[list[str], list[str]]:
    """Check that the Dockerfile follows base-image conventions.

    Returns (issues, warnings).  Issues are hard failures; warnings are
    acceptable temporary states (e.g. full install for an env not yet in base).
    """
    dockerfile = task_dir / "environment" / "Dockerfile"
    if not dockerfile.exists():
        return [], []  # already caught by format check

    src = dockerfile.read_text()
    issues: list[str] = []
    warnings: list[str] = []

    # Must use the shared base image
    if not BASE_IMAGE_PATTERN.search(src):
        issues.append(
            "Dockerfile: not using smolclaws-base image "
            "(should be FROM kywch/smolclaws-base:latest)"
        )

    # Must create log directories
    if "/logs/verifier" not in src:
        issues.append(
            "Dockerfile: missing log directory creation "
            "(need: mkdir -p /logs/verifier /logs/agent /logs/artifacts)"
        )

    # Should not hardcode GOOGLE_WORKSPACE_CLI_API_BASE_URL — use gws wrapper
    if "GOOGLE_WORKSPACE_CLI_API_BASE_URL" in src:
        issues.append(
            "Dockerfile: hardcoded GOOGLE_WORKSPACE_CLI_API_BASE_URL "
            "(use CLAW_*_URL env vars + gws wrapper instead)"
        )

    # Should use uv pip install --no-deps (deps are in base).
    # Exception: if the env is not yet in the base image, a full install
    # is acceptable for testing — emit a warning instead of a failure.
    if "pip install" in src and "--no-deps" not in src:
        env_pkgs = _dockerfile_env_packages(src)
        base_envs = _envs_in_base_image()
        not_in_base = env_pkgs - base_envs
        if not_in_base:
            names = ", ".join(sorted(not_in_base))
            warnings.append(
                f"Dockerfile: {names} not in base image — "
                f"full install OK for testing, switch to --no-deps after base update"
            )
        else:
            issues.append(
                "Dockerfile: pip install without --no-deps "
                "(deps are pre-installed in base image, use: "
                "uv pip install --system --no-deps -e .)"
            )

    # Must set TASKS_DIR=/tasks for evaluate.py to find needles.py
    if "TASKS_DIR" not in src:
        issues.append(
            "Dockerfile: missing ENV TASKS_DIR=/tasks "
            "(required for evaluate.py to find needles.py in the container)"
        )

    # Must install to /app-<env>, not /app (sandbox convention)
    if re.search(r"COPY\s+packages/environments/claw-\w+\s+/app\b(?!-)", src):
        issues.append(
            "Dockerfile: installing env to /app instead of /app-<env> "
            "(use /app-gdrive, /app-gmail, etc. — /app is locked by base image)"
        )

    # Must lock env source code and grant agent log access
    if "chmod 700" not in src:
        issues.append(
            "Dockerfile: missing chmod 700 /app-<env> "
            "(must lock env source code to prevent agent from reading it)"
        )
    if "chown agent:agent" not in src:
        issues.append(
            "Dockerfile: missing chown agent:agent /logs/agent /logs/artifacts "
            "(must grant agent write access to log directories)"
        )

    return issues, warnings


# ---------------------------------------------------------------------------
# Skills duplication check
# ---------------------------------------------------------------------------

def check_no_inline_skills(task_dir: Path) -> list[str]:
    """Flag tasks that embed skills inline instead of using top-level skills/."""
    skills_dir = task_dir / "environment" / "skills"
    if skills_dir.is_dir():
        skill_names = [d.name for d in skills_dir.iterdir() if d.is_dir()]
        return [
            f"environment/skills/: inline skills found ({', '.join(skill_names)}) "
            f"— skills should live in top-level skills/ directory"
        ]
    return []


# ---------------------------------------------------------------------------
# Legacy docker-compose check
# ---------------------------------------------------------------------------

def check_no_docker_compose(task_dir: Path) -> list[str]:
    """Flag tasks that ship a docker-compose.yaml alongside a base-image Dockerfile.

    The base image runs environments in-process via CLAW_*_URL on localhost,
    so a separate compose service is unnecessary.  A docker-compose.yaml next
    to a base-image Dockerfile is a sign of an incomplete migration.
    """
    compose = task_dir / "environment" / "docker-compose.yaml"
    if not compose.exists():
        compose = task_dir / "environment" / "docker-compose.yml"
    if not compose.exists():
        return []

    # Only flag if the Dockerfile already uses the base image
    dockerfile = task_dir / "environment" / "Dockerfile"
    if dockerfile.exists() and BASE_IMAGE_PATTERN.search(dockerfile.read_text()):
        return [
            "environment/docker-compose.yaml: legacy docker-compose found "
            "— base image runs environments in-process via CLAW_*_URL, "
            "remove docker-compose.yaml"
        ]
    return []


# ---------------------------------------------------------------------------
# task.toml schema check
# ---------------------------------------------------------------------------

TASK_TOML_KNOWN_SECTIONS = {"metadata", "agent", "verifier", "environment"}
TASK_TOML_KNOWN_TOP = {"version"}
TASK_TOML_ENV_KNOWN = {
    "cpus", "memory_mb", "allow_internet", "skills_dir",
    "docker_compose", "ports",
}


def check_task_toml_schema(task_dir: Path) -> list[str]:
    """Validate task.toml has expected structure and no unknown fields."""
    toml_path = task_dir / "task.toml"
    if not toml_path.exists():
        return []  # already caught by format check

    try:
        data = tomllib.loads(toml_path.read_text())
    except Exception as exc:
        return [f"task.toml: parse error — {exc}"]

    issues: list[str] = []

    # Check top-level keys
    for key in data:
        if key not in TASK_TOML_KNOWN_SECTIONS and key not in TASK_TOML_KNOWN_TOP:
            issues.append(f"task.toml: unknown top-level key '{key}'")

    # Check required top-level
    if "version" not in data:
        issues.append("task.toml: missing 'version' field")

    # Check environment section
    env = data.get("environment", {})
    for key in env:
        if key not in TASK_TOML_ENV_KNOWN:
            issues.append(f"task.toml: unknown environment key '{key}'")

    return issues


# ---------------------------------------------------------------------------
# config.toml cross-check
# ---------------------------------------------------------------------------

def _load_environments_toml() -> dict:
    """Load the canonical config.toml from repo root."""
    path = REPO_ROOT / "config.toml"
    if not path.exists():
        return {}
    return tomllib.loads(path.read_text())


def check_base_image_consistency() -> list[str]:
    """Check that Dockerfile.base includes all envs from config.toml."""
    envs = _load_environments_toml()
    if not envs:
        return []

    dockerfile_base = REPO_ROOT / "docker" / "Dockerfile.base"
    if not dockerfile_base.exists():
        return ["docker/Dockerfile.base not found"]

    src = dockerfile_base.read_text()
    issues: list[str] = []

    for env_name, env_cfg in envs.items():
        if not env_name.startswith("claw-"):
            continue
        # Check package is copied into base image deps
        if f"packages/environments/{env_name}" not in src:
            issues.append(
                f"Dockerfile.base: missing dependency install for {env_name} "
                f"(defined in config.toml but not in base image)"
            )

        # Check that a CLAW_*_URL env var is set for this service's port.
        # Naming varies (claw-gcal -> CLAW_CALENDAR_URL), so match by port.
        port = env_cfg.get("port")
        if port and f"localhost:{port}" not in src:
            issues.append(
                f"Dockerfile.base: no CLAW_*_URL env var for {env_name} "
                f"(port {port} defined in config.toml)"
            )

    return issues


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> int:
    parser = argparse.ArgumentParser(description="Validate (and optionally fix) task directories.")
    parser.add_argument("--task", metavar="NAME", help="Check a single task by name.")
    parser.add_argument("--skip-verified", action="store_true",
                        help="Skip tasks that already have oracle_score=1.0 and unchanged files.")
    parser.add_argument("--check-base", action="store_true",
                        help="Also check Dockerfile.base against config.toml.")
    args = parser.parse_args()

    if not TASKS_DIR.exists():
        print(f"ERROR: tasks dir not found: {TASKS_DIR}", file=sys.stderr)
        return 2

    cache = {}
    if args.skip_verified:
        from clawbench.score_cache import load_cache, is_verified
        cache = load_cache(TASKS_DIR)

    if args.task:
        task_dirs = [TASKS_DIR / args.task]
    else:
        task_dirs = sorted(
            d for d in TASKS_DIR.iterdir()
            if d.is_dir()
            and d.name not in SKIP
            and not d.name.startswith((".", "_"))
        )

    total = 0
    failures = 0
    warned = 0

    for task_dir in task_dirs:
        task = task_dir.name
        if args.skip_verified and cache:
            from clawbench.score_cache import is_verified
            if is_verified(task, task_dir, cache):
                print(f"SKIP {task}  (verified, oracle=1.0)")
                continue
        issues: list[str] = []
        task_warnings: list[str] = []

        # Harbor format checks
        issues.extend(check_format(task_dir))

        # evaluate.py checks
        issues.extend(check_evaluate_callable(task_dir))
        issues.extend(check_rewards_contract(task_dir))

        # Dockerfile convention checks
        df_issues, df_warnings = check_dockerfile_conventions(task_dir)
        issues.extend(df_issues)
        task_warnings.extend(df_warnings)

        # Inline skills check
        issues.extend(check_no_inline_skills(task_dir))

        # Legacy docker-compose check
        issues.extend(check_no_docker_compose(task_dir))

        # task.toml schema check
        issues.extend(check_task_toml_schema(task_dir))

        total += 1
        if issues:
            failures += 1
            for issue in issues:
                print(f"FAIL {task}: {issue}")
        else:
            for w in task_warnings:
                print(f"WARN {task}: {w}")
            if task_warnings:
                warned += 1
                print(f"OK   {task} ({len(task_warnings)} warning(s))")
            else:
                print(f"OK   {task}")

    # Cross-cutting: Dockerfile.base vs config.toml
    if args.check_base:
        base_issues = check_base_image_consistency()
        if base_issues:
            failures += 1
            for issue in base_issues:
                print(f"FAIL [base-image]: {issue}")
        else:
            print("OK   [base-image]")

    summary = f"\n{total - failures}/{total} tasks passed."
    if warned:
        summary += f" ({warned} with warnings)"
    print(summary)
    return 0 if failures == 0 else 1


if __name__ == "__main__":
    sys.exit(main())
