"""Stage 1: LLM-powered source code analysis of reward functions."""

import inspect
import json
import re
from dataclasses import dataclass, field
from typing import Any

from rewardprobe.tier2.client import LLMClient
from rewardprobe.tier2.prompts import CODE_ANALYSIS_SYSTEM, CODE_ANALYSIS_USER
from rewardprobe.types import NormalizedRewardFn


@dataclass
class ExpectedBehavior:
    should_reject_empty: bool = True
    should_reject_random_text: bool = True
    checks_answer_correctness: bool = True
    checks_output_format: bool = False


@dataclass
class Vulnerability:
    name: str
    severity: str  # "critical" or "warning"
    description: str
    exploit: str
    expected_score: float = 0.0
    actual_score_estimate: float = 1.0


@dataclass
class CodeAnalysis:
    function_name: str
    function_type: str  # correctness, format, auxiliary, judge
    what_it_checks: str
    expected_behavior: ExpectedBehavior
    vulnerabilities: list[Vulnerability] = field(default_factory=list)
    source_code: str = ""
    raw_response: dict = field(default_factory=dict)


def _get_source_code(fn: NormalizedRewardFn) -> str:
    """Try to get the source code of a reward function."""
    # Try the wrapped fn first
    target = fn.fn
    # If it's a wrapper, try to find the original
    if hasattr(target, "__wrapped__"):
        target = target.__wrapped__
    # Try closure variables for our adapter wrappers
    if hasattr(target, "__closure__") and target.__closure__:
        for cell in target.__closure__:
            try:
                val = cell.cell_contents
                if callable(val) and hasattr(val, "__name__"):
                    try:
                        return inspect.getsource(val)
                    except (OSError, TypeError):
                        pass
            except ValueError:
                pass
    try:
        return inspect.getsource(target)
    except (OSError, TypeError):
        return f"# Source code not available for {fn.name}"


def _get_dependency_sources(source: str) -> str:
    """Try to find source code of functions called by the reward function."""
    # Look for function calls in the source and try to get their source
    # This is best-effort — we can't always resolve all dependencies
    re.findall(r"(\w+)\s*\(", source)
    # We can't reliably resolve these without the module context,
    # but we include what we find via the closure/globals
    return ""


def analyze_reward_fn(
    fn: NormalizedRewardFn,
    tasks: list[dict[str, Any]],
    client: LLMClient,
) -> CodeAnalysis:
    """Analyze a reward function using an LLM."""
    source = _get_source_code(fn)
    deps = _get_dependency_sources(source)

    # Format sample tasks (first 3, truncated)
    sample = json.dumps(tasks[:3], indent=2, default=str)[:1000]

    deps_section = ""
    if deps:
        deps_section = f"\nHelper functions used:\n```python\n{deps}\n```"

    prompt = CODE_ANALYSIS_USER.format(
        source_code=source,
        dependencies_code=deps_section,
        sample_tasks=sample,
    )

    try:
        result = client.complete_json(CODE_ANALYSIS_SYSTEM, prompt)
    except (json.JSONDecodeError, Exception) as e:
        # Fallback: return a conservative analysis
        return CodeAnalysis(
            function_name=fn.name,
            function_type="unknown",
            what_it_checks=f"Could not analyze: {e}",
            expected_behavior=ExpectedBehavior(),
            source_code=source,
        )

    # Parse response
    eb = result.get("expected_behavior", {})
    expected = ExpectedBehavior(
        should_reject_empty=eb.get("should_reject_empty", True),
        should_reject_random_text=eb.get("should_reject_random_text", True),
        checks_answer_correctness=eb.get("checks_answer_correctness", True),
        checks_output_format=eb.get("checks_output_format", False),
    )

    vulns = []
    for v in result.get("vulnerabilities", []):
        vulns.append(
            Vulnerability(
                name=v.get("name", "unknown"),
                severity=v.get("severity", "warning"),
                description=v.get("description", ""),
                exploit=v.get("exploit", ""),
                expected_score=v.get("expected_score", 0.0),
                actual_score_estimate=v.get("actual_score_estimate", 1.0),
            )
        )

    return CodeAnalysis(
        function_name=fn.name,
        function_type=result.get("function_type", "unknown"),
        what_it_checks=result.get("what_it_checks", "unknown"),
        expected_behavior=expected,
        vulnerabilities=vulns,
        source_code=source,
        raw_response=result,
    )
