"""Click CLI for rewardprobe."""

from __future__ import annotations

import json
import sys
from pathlib import Path

import click
from rich.console import Console

from rewardprobe.probe import Probe
from rewardprobe.report import print_report, report_as_json, report_to_json

console = Console(stderr=True)


def _load_jsonl(path: str) -> list[dict]:
    """Load a JSONL file — one JSON object per line."""
    tasks: list[dict] = []
    p = Path(path)
    if not p.is_file():
        raise click.BadParameter(f"File not found: {path}", param_hint="'--dataset'")
    for i, line in enumerate(p.read_text().splitlines(), 1):
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError as exc:
            raise click.BadParameter(
                f"Invalid JSON on line {i} of {path}: {exc}",
                param_hint="'--dataset'",
            ) from exc
        if not isinstance(obj, dict):
            raise click.BadParameter(
                f"Line {i} of {path} is not a JSON object (got {type(obj).__name__})",
                param_hint="'--dataset'",
            )
        tasks.append(obj)
    if not tasks:
        raise click.BadParameter(f"No tasks found in {path}", param_hint="'--dataset'")
    return tasks


def _resolve_target(target: str, dataset: str | None, adapter: str | None):
    """Determine how to load the target and return (mode, details).

    Returns one of:
        ("raw", module_path, function_name, tasks)
        ("verifiers", file_path)
        ("trl", module_path, function_name, tasks)

    Raises click.UsageError on problems.
    """
    # If target is a .jsonl file, that's an error
    if target.endswith(".jsonl"):
        raise click.UsageError(
            f"'{target}' looks like a dataset file, not a target.\n"
            f"Did you mean: rewardprobe test <your_module> --dataset {target}"
        )

    # Explicit adapter override
    if adapter == "trl":
        if "::" not in target:
            raise click.UsageError(
                "TRL adapter requires module::function format.\n"
                "Example: rewardprobe test my_rewards::compute_reward "
                "--adapter trl --dataset tasks.jsonl"
            )
        module_path, fn_name = target.split("::", 1)
        tasks = _require_dataset(dataset, "TRL")
        return ("trl", module_path, fn_name, tasks)

    # module::function → raw adapter
    if "::" in target:
        module_path, fn_name = target.split("::", 1)
        if adapter and adapter == "verifiers":
            raise click.UsageError(
                "Cannot use verifiers adapter with module::function format.\n"
                "Pass a .py file path instead."
            )
        tasks = _require_dataset(dataset, "raw function")
        mode = adapter or "raw"
        return (mode, module_path, fn_name, tasks)

    # .py file — check for load_environment()
    target_path = Path(target)
    if target_path.suffix == ".py" and target_path.is_file():
        if adapter == "raw":
            raise click.UsageError(
                "Raw adapter requires module::function format.\n"
                "Example: rewardprobe test my_module::my_fn --dataset tasks.jsonl"
            )
        # Check if it has load_environment
        if adapter == "verifiers" or _file_has_load_environment(target_path):
            return ("verifiers", str(target_path.resolve()))
        # Fall back: try to detect what's in the file
        raise click.UsageError(
            f"'{target}' does not contain load_environment().\n"
            f"If this is a module with reward functions, use: "
            f"rewardprobe test {target_path.stem}::function_name --dataset tasks.jsonl"
        )

    # Try as importable module
    if adapter == "verifiers":
        raise click.UsageError(f"Verifiers adapter expects a .py file path, got '{target}'.")

    # Last resort: try importing as a module path
    try:
        _try_import(target)
        raise click.UsageError(
            f"'{target}' is a module but no function was specified.\n"
            f"Use module::function format: "
            f"rewardprobe test {target}::your_function --dataset tasks.jsonl"
        )
    except ImportError:
        pass

    raise click.UsageError(
        f"Cannot determine how to test '{target}'.\n\n"
        f"Expected one of:\n"
        f"  rewardprobe test module::function --dataset tasks.jsonl   (raw function)\n"
        f"  rewardprobe test environment.py                           (verifiers environment)\n"
        f"  rewardprobe test module::function --adapter trl --dataset tasks.jsonl  (TRL)"
    )


def _require_dataset(dataset: str | None, adapter_name: str) -> list[dict]:
    """Ensure --dataset was provided and load it."""
    if not dataset:
        raise click.UsageError(
            f"{adapter_name} testing requires a dataset.\nAdd: --dataset tasks.jsonl"
        )
    return _load_jsonl(dataset)


def _file_has_load_environment(path: Path) -> bool:
    """Check if a Python file defines load_environment using AST parsing."""
    import ast

    try:
        source = path.read_text()
        tree = ast.parse(source)
    except (OSError, SyntaxError):
        return False
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef) and node.name == "load_environment":
            return True
    return False


def _try_import(module_path: str):
    """Try to import a module by dotted path."""
    import importlib

    importlib.import_module(module_path)


def _import_function(module_path: str, fn_name: str):
    """Import a function from a module path (dotted or file path)."""
    import importlib
    import importlib.util

    # Try as file path first
    p = Path(module_path)
    if p.suffix == ".py" and p.is_file():
        spec = importlib.util.spec_from_file_location("_target_module", str(p.resolve()))
        if spec is None or spec.loader is None:
            raise click.UsageError(f"Cannot load module from '{module_path}'")
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)  # type: ignore[union-attr]
    else:
        # Try as dotted module path
        try:
            mod = importlib.import_module(module_path)
        except ImportError as exc:
            raise click.UsageError(f"Cannot import module '{module_path}': {exc}") from exc

    if not hasattr(mod, fn_name):
        available = [a for a in dir(mod) if not a.startswith("_") and callable(getattr(mod, a))]
        hint = ""
        if available:
            hint = f"\nAvailable callables: {', '.join(available[:10])}"
        raise click.UsageError(f"Module '{module_path}' has no attribute '{fn_name}'.{hint}")
    fn = getattr(mod, fn_name)
    if not callable(fn):
        raise click.UsageError(f"'{module_path}::{fn_name}' is not callable")
    return fn


# ── CLI definition ───────────────────────────────────────────────────


@click.group(invoke_without_command=True)
@click.version_option(package_name="rewardprobe", prog_name="rewardprobe")
@click.pass_context
def main(ctx: click.Context):
    """rewardprobe -- Pre-training stress-testing for reward functions."""
    if ctx.invoked_subcommand is None:
        click.echo(ctx.get_help())


@main.command()
@click.argument("target")
@click.option("--dataset", "-d", default=None, help="Path to JSONL tasks file (for raw functions).")
@click.option(
    "--adapter",
    "-a",
    type=click.Choice(["verifiers", "trl", "raw"], case_sensitive=False),
    default=None,
    help="Force adapter type. Auto-detected if omitted.",
)
@click.option("--ci", is_flag=True, help="Exit with code 1 on any critical finding.")
@click.option("--strict", is_flag=True, help="Exit 1 on any finding (incl. warnings).")
@click.option("--num-tasks", "-n", default=10, show_default=True, help="Tasks per attack.")
@click.option(
    "--families",
    "-f",
    default=None,
    help="Comma-separated family names to run (default: all).",
)
@click.option("--json", "json_output", is_flag=True, help="Output JSON instead of terminal report.")
@click.option("--output", "-o", default=None, help="Write JSON report to file.")
@click.option("--seed", default=42, show_default=True, help="Random seed.")
def test(
    target: str,
    dataset: str | None,
    adapter: str | None,
    ci: bool,
    strict: bool,
    num_tasks: int,
    families: str | None,
    json_output: bool,
    output: str | None,
    seed: int,
):
    """Test a reward function or verifiers environment.

    TARGET is a path to a .py environment file or a module::function reference.

    \b
    Examples:
      rewardprobe test my_env.py
      rewardprobe test rewards::check_answer --dataset tasks.jsonl
      rewardprobe test rewards::score --adapter trl --dataset tasks.jsonl
    """
    family_list = [f.strip() for f in families.split(",")] if families else None

    probe = Probe(
        families=family_list,
        num_tasks=num_tasks,
        seed=seed,
    )

    try:
        resolved = _resolve_target(target, dataset, adapter)
    except click.UsageError:
        raise
    except Exception as exc:
        if ci or strict:
            console.print(f"[red]Error:[/red] {exc}")
            sys.exit(1)
        raise

    try:
        mode = resolved[0]
        if mode == "verifiers":
            report = probe.test_environment(resolved[1])
        elif mode == "trl":
            _, module_path, fn_name, tasks = resolved
            fn = _import_function(module_path, fn_name)
            report = probe.test_trl(fn, tasks)
        else:
            # raw
            _, module_path, fn_name, tasks = resolved
            fn = _import_function(module_path, fn_name)
            report = probe.test_fn(fn, tasks, name=fn_name)
    except click.UsageError:
        raise
    except Exception as exc:
        if ci or strict:
            console.print(f"[red]Error:[/red] {exc}")
            sys.exit(1)
        raise click.ClickException(str(exc))

    # Output
    if output:
        report_to_json(report, output)
        if not json_output:
            console.print(f"[dim]Report written to {output}[/dim]")

    if json_output:
        out_console = Console()
        out_console.print(report_as_json(report), highlight=False)
    else:
        out_console = Console()
        print_report(report, console=out_console)

    # Exit codes
    if ci and report.critical_count > 0:
        sys.exit(1)
    if strict and len(report.findings) > 0:
        sys.exit(1)
