"""Probe — the main orchestrator and Python API entry point."""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Callable

from rewardprobe.types import (
    FamilyResult,
    Finding,
    NormalizedEnvironment,
    NormalizedRewardFn,
    Report,
    Severity,
)


class Probe:
    """Stress-test reward functions before training.

    The main entry point for rewardprobe. Normalizes any reward function
    format into NormalizedEnvironment, then runs selected attack families.
    """

    def __init__(
        self,
        families: list[str] | None = None,
        num_tasks: int = 10,
        seed: int = 42,
        deep: bool = False,
        model: str = "opus",
    ):
        """
        Args:
            families: which attack families to run (None = all).
            num_tasks: how many tasks to sample per attack.
            seed: random seed for reproducibility.
            deep: if True, run Tier 2 LLM-powered analysis after Tier 1.
            model: LLM model for Tier 2 (haiku, sonnet, opus, gpt-4o, etc).
        """
        self.families = families
        self.num_tasks = num_tasks
        self.seed = seed
        self.deep = deep
        self.model = model

    # ------------------------------------------------------------------
    # Public API — one method per adapter style
    # ------------------------------------------------------------------

    def test_fn(
        self,
        fn: Callable | list[Callable],
        tasks: list[dict[str, Any]],
        name: str | None = None,
        weights: list[float] | None = None,
    ) -> Report:
        """Test any reward function(s). Auto-detects signature pattern.

        Works with raw Python functions, TRL batch functions,
        and any other callable — no framework-specific flags needed.
        """
        from rewardprobe.adapters.auto import auto_adapt

        env = auto_adapt(fn, tasks, name=name, weights=weights)
        return self._run_attacks(env)

    def simulate_fn(
        self,
        fn: Callable | list[Callable],
        tasks: list[dict[str, Any]],
        name: str | None = None,
        weights: list[float] | None = None,
        n_tasks: int = 5,
    ) -> Any:
        """Simulate the reward landscape. Requires an LLM API key.

        Generates diverse completions, scores them, and analyzes
        what strategies a model would learn during training.
        """
        from rewardprobe.adapters.auto import auto_adapt
        from rewardprobe.simulator import simulate
        from rewardprobe.tier2.client import get_client

        env = auto_adapt(fn, tasks, name=name, weights=weights)
        client = get_client(self.model)
        return simulate(env, client, n_tasks=n_tasks)

    def test_trl(
        self,
        fn: Callable,
        tasks: list[dict[str, Any]],
        format: str = "standard",
        columns: list[str] | None = None,
    ) -> Report:
        """Test TRL-style reward functions. Deprecated — use test_fn() instead."""
        return self.test_fn(fn, tasks)

    def test_environment(self, path: str) -> Report:
        """Test a verifiers environment file.

        Imports the module at *path*, calls ``load_environment()``, and
        uses VerifiersAdapter to produce a NormalizedEnvironment.
        """
        import importlib.util
        import os

        abs_path = os.path.abspath(path)
        if not os.path.isfile(abs_path):
            raise FileNotFoundError(f"Environment file not found: {abs_path}")

        spec = importlib.util.spec_from_file_location("_env_module", abs_path)
        if spec is None or spec.loader is None:
            raise ImportError(f"Cannot load module from {abs_path}")
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)  # type: ignore[union-attr]

        if not hasattr(module, "load_environment"):
            raise AttributeError(
                f"{abs_path} does not define load_environment(). "
                "Is this a verifiers environment file?"
            )

        env_obj = module.load_environment()

        # Try to use VerifiersAdapter if available; fall back to duck-typing
        try:
            from rewardprobe.adapters.verifiers import VerifiersAdapter

            adapter = VerifiersAdapter()
            env = adapter.normalize(env_obj)
        except ImportError:
            # Minimal duck-type extraction when the verifiers adapter
            # hasn't been written yet — keeps the orchestrator testable.
            env = self._duck_type_environment(env_obj, path)

        return self._run_attacks(env)

    # ------------------------------------------------------------------
    # Core runner
    # ------------------------------------------------------------------

    def _run_attacks(self, env: NormalizedEnvironment) -> Report:
        """Run selected attack families against a normalized environment."""
        from rewardprobe.families import FAMILIES, FAMILY_ORDER
        from rewardprobe.families.base import _normalize_answer_key

        # Normalize answer keys before sampling
        normalized_tasks = _normalize_answer_key(env.tasks)
        env_normalized = env.model_copy(update={"tasks": normalized_tasks})

        # Sample tasks for reproducibility
        sampled_tasks = env_normalized.sample_tasks(n=self.num_tasks, seed=self.seed)
        env_sampled = env_normalized.model_copy(update={"tasks": sampled_tasks})

        # Determine which families to run
        if self.families is not None:
            family_names = [f.lower() for f in self.families]
            ordered = [name for name in FAMILY_ORDER if name in family_names]
            # Warn about unknown family names
            known = set(FAMILY_ORDER)
            for name in family_names:
                if name not in known:
                    raise ValueError(
                        f"Unknown family '{name}'. Available: {', '.join(FAMILY_ORDER)}"
                    )
        else:
            ordered = list(FAMILY_ORDER)

        # Run each family
        results: list[FamilyResult] = []
        for family_name in ordered:
            family_cls = FAMILIES.get(family_name)
            if family_cls is None:
                continue
            family_instance = family_cls()
            result = family_instance.run(env_sampled)
            results.append(result)

        # ── Tier 2: LLM-powered deep analysis ──
        tier2_results: list[FamilyResult] = []

        if self.deep:
            from rewardprobe.tier2.adversarial import (
                adversarial_to_findings,
                generate_adversarial,
            )
            from rewardprobe.tier2.client import get_client
            from rewardprobe.tier2.code_analyzer import analyze_reward_fn
            from rewardprobe.tier2.validator import filter_tier1_findings

            client = get_client(self.model)

            # Stage 1: Analyze each reward function's source code
            analyses = {}
            for fn in env_sampled.reward_fns:
                analysis = analyze_reward_fn(fn, env_sampled.tasks, client)
                analyses[fn.name] = analysis

            # Filter Tier 1 false positives using analysis
            all_t1_findings = [f for r in results for f in r.findings]
            kept, removed = filter_tier1_findings(all_t1_findings, analyses)
            len(removed)

            # Rebuild Tier 1 results with filtered findings
            if removed:
                removed_set = set(id(f) for f in removed)
                for result in results:
                    result.findings = [f for f in result.findings if id(f) not in removed_set]
                    result.passed = result.total - len({f.attack for f in result.findings})

            # Stage 1 findings from code analysis
            code_findings: list[Finding] = []
            for fn_name, analysis in analyses.items():
                for vuln in analysis.vulnerabilities:
                    # Skip speculative findings about unresolved dependencies
                    desc_lower = vuln.description.lower()
                    if any(
                        kw in desc_lower
                        for kw in [
                            "not provided",
                            "not defined",
                            "missing",
                            "cannot verify",
                            "without seeing",
                            "not available",
                        ]
                    ):
                        continue

                    # Verify the exploit actually works
                    fn_obj = next(
                        (f for f in env_sampled.reward_fns if f.name == fn_name),
                        None,
                    )
                    actual_score = None
                    if fn_obj and vuln.exploit:
                        from rewardprobe.families.base import _safe_call

                        task = env_sampled.tasks[0] if env_sampled.tasks else {}
                        score, _ = _safe_call(fn_obj, vuln.exploit, task)
                        actual_score = score

                    # Only include findings where exploit was verified OR
                    # is a structural concern (no exploit needed)
                    if vuln.exploit and actual_score is not None and actual_score < 0.1:
                        continue  # Exploit didn't work — skip

                    sev = Severity.CRITICAL if vuln.severity == "critical" else Severity.WARNING
                    code_findings.append(
                        Finding(
                            attack="code_analysis",
                            family="ai_analysis",
                            severity=sev,
                            input=vuln.exploit[:300] if vuln.exploit else "",
                            actual_reward=actual_score,
                            recommendation=(f"[AI] {vuln.name}: {vuln.description}"),
                        )
                    )

            if code_findings:
                from rewardprobe.families.base import _aggregate_findings

                tier2_results.append(
                    FamilyResult(
                        name="ai_analysis",
                        display_name="AI Code Analysis",
                        total=len(analyses),
                        passed=len(analyses) - len(code_findings),
                        findings=_aggregate_findings(code_findings),
                    )
                )

            # Stage 2: Generate adversarial completions
            adv_findings: list[Finding] = []
            correctness_fns = [
                (fn, analyses[fn.name])
                for fn in env_sampled.reward_fns
                if fn.name in analyses
                and analyses[fn.name].expected_behavior.checks_answer_correctness
            ]

            if correctness_fns:
                # Test on a subset of tasks
                adv_tasks = env_sampled.tasks[: min(3, len(env_sampled.tasks))]
                for fn, analysis in correctness_fns:
                    for idx, task in enumerate(adv_tasks):
                        adv_results = generate_adversarial(fn, task, analysis, client)
                        adv_findings.extend(adversarial_to_findings(adv_results, fn.name, idx))

            if adv_findings:
                from rewardprobe.families.base import _aggregate_findings

                tier2_results.append(
                    FamilyResult(
                        name="adversarial",
                        display_name="AI Adversarial",
                        total=len(correctness_fns) * min(3, len(env_sampled.tasks)),
                        passed=len(correctness_fns) * min(3, len(env_sampled.tasks))
                        - len(adv_findings),
                        findings=_aggregate_findings(adv_findings),
                    )
                )

        # Build report
        all_results = results + tier2_results
        all_passed = all(r.all_passed for r in all_results)
        version = _get_version()

        return Report(
            environment=env.name,
            rewardprobe_version=version,
            timestamp=datetime.now(timezone.utc).isoformat(),
            passed=all_passed,
            families=all_results,
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _duck_type_environment(env_obj: Any, path: str) -> NormalizedEnvironment:
        """Best-effort extraction from an unknown environment object."""
        import os

        name = getattr(env_obj, "name", os.path.basename(path))
        tasks = getattr(env_obj, "tasks", [])
        reward_fns_raw = getattr(env_obj, "reward_fns", [])

        reward_fns: list[NormalizedRewardFn] = []
        for i, fn in enumerate(reward_fns_raw):
            if isinstance(fn, NormalizedRewardFn):
                reward_fns.append(fn)
            elif callable(fn):
                reward_fns.append(
                    NormalizedRewardFn(
                        name=getattr(fn, "__name__", f"reward_fn_{i}"),
                        fn=fn,
                    )
                )

        if not reward_fns:
            raise ValueError(
                f"Could not extract reward functions from {path}. "
                "Make sure load_environment() returns an object with "
                "a 'reward_fns' attribute."
            )

        return NormalizedEnvironment(
            name=name,
            reward_fns=reward_fns,
            tasks=list(tasks) if tasks else [],
            metadata={"adapter": "verifiers", "source": path},
        )


def _get_version() -> str:
    """Return rewardprobe version string."""
    try:
        from rewardprobe import __version__

        return __version__
    except ImportError:
        return "0.1.0"
