""" Agents for Task3 : Rule Checking for a function """

import json
import random as _random
from typing import Any, Dict, List

from server import Task3Environment
from env.schemas import Action, ActionType
from data.data_loader import load_contracts, get_function_by_name

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _parse_fn_list(result_text: str) -> List[str]:
    """Parse 'Functions in X: f1, f2, f3' into [f1, f2, f3]."""
    if ": " in result_text:
        return [f.strip() for f in result_text.split(": ", 1)[-1].split(", ") if f.strip()]
    return []


# ─────────────────────────────────────────────────────────────────────────────
# Task 3 agents
# ─────────────────────────────────────────────────────────────────────────────

def oracle_t3(env: Task3Environment, seed: int, verbose: bool = False) -> Dict[str, Any]:
    """Submits exact target function → score = 1.0."""
    r        = env.reset(seed=seed)
    obs      = r.observation
    fn_name  = env.state().target_function
    contract = obs.contract_name
    if verbose:
        prop = obs.extra.get("property_english", "")[:60]
        print(f"    {contract}.{fn_name}()  \"{prop}\"")

    env.step(Action(action_type=ActionType.GET_PROPERTY_SPECIFICATION))
    env.step(Action(action_type=ActionType.LIST_FUNCTIONS))
    result = env.step(Action(action_type=ActionType.SUBMIT_FUNCTION,
                              params={"function_name": fn_name}))
    return {"seed": seed, "contract": contract, "target_function": fn_name,
            "grader_score": result.reward.value}


def subfunction_t3(env: Task3Environment, seed: int, verbose: bool = False) -> Dict[str, Any]:
    """Submits the first partial-credit subfunction if one exists, else 'constructor'."""
    r   = env.reset(seed=seed)
    obs = r.observation
    contracts = load_contracts()
    target_contract = {}
    for c in contracts:
        if c["contract_name"] == obs.contract_name:
            target_contract = c
            break
    
    submit_name = "constructor"
    if (target_contract and "call_graph" in target_contract and env.state().target_function in 
        target_contract["call_graph"] and target_contract["call_graph"][env.state().target_function]):
        submit_name = target_contract["call_graph"][env.state().target_function][0]
    
        
    result = env.step(Action(action_type=ActionType.SUBMIT_FUNCTION,
                              params={"function_name": submit_name}))
    if verbose:
        prop = obs.extra.get("property_english", "")[:60]
        print(f"    {obs.contract_name}.{env.state().target_function}()  \"{prop}\"")
        print(f"    Submitting subfunction: {submit_name}")
        print(f"    Reward received: {result.reward.value}")
        
    return {"seed": seed, "grader_score": result.reward.value, "submitted": submit_name}


def random_t3(env: Task3Environment, seed: int) -> Dict[str, Any]:
    """Genuine random agent: lists functions, picks one at random, submits.

    With N functions per contract and 1 target, expected score ≈ 1/N ≈ 0.20–0.25.
    Uses a seeded RNG for reproducibility.
    """
    rng = _random.Random(seed ^ 0xCAFE1)
    env.reset(seed=seed)

    # Step 1: get function list (necessary to pick a real candidate)
    s   = env.step(Action(action_type=ActionType.LIST_FUNCTIONS))
    fns = _parse_fn_list(s.observation.last_action_result or "")
    if not fns:
        fns = ["constructor"]

    # Step 2: optionally do 1 cheap browse action (formalized or call_graph)
    browse_options = [
        (ActionType.GET_PROPERTY_SPECIFICATION, {}),
        (ActionType.GET_CALL_GRAPH,          {}),
    ]
    at, params = rng.choice(browse_options)
    env.step(Action(action_type=at, params=params))

    # Step 3: submit a uniformly random function from the real list
    chosen = rng.choice(fns)
    result = env.step(Action(action_type=ActionType.SUBMIT_FUNCTION,
                              params={"function_name": chosen}))

    return {"seed": seed, "grader_score": result.reward.value, "submitted": chosen}


def floor_t3(env: Task3Environment, seed: int) -> Dict[str, Any]:
    """Always submits 'constructor' → guaranteed score = 0.0."""
    env.reset(seed=seed)
    result = env.step(Action(action_type=ActionType.SUBMIT_FUNCTION,
                              params={"function_name": "constructor"}))
    return {"seed": seed, "grader_score": 0.001}