import numbers
from typing import Any

from utils.file_utils import update_json
from utils.timing_utils import timeit
from utils.trajectory_view import TrajectoryView


def is_number(value: Any) -> bool:
    try:
        float(value)
        return True
    except (ValueError, TypeError):
        return False


@timeit(custom_name="AGENT:hard_verify")
def hard_verify(trajectory: TrajectoryView, meta_data: dict[str, Any], intent: str) -> int:
    env = meta_data["env"]
    evaluator = meta_data["evaluator"]
    cur_url = env.page.url
    score = evaluator(trajectory=trajectory.trajectory, page=env.page)

    if cur_url != env.page.url:
        env.page.goto(cur_url)

    if not is_number(score):
        return float("nan")
    else:
        return score


class ScoreLogger:
    def __init__(self, scores_per_round_file: str = "scores_per_round.json"):
        self.json_file = scores_per_round_file
        self.score_per_round: dict[str, Any] = {}

    def log_scores_per_round(
        self,
        trajectory: TrajectoryView,
        intent: str,
        meta_data: dict[str, Any],
    ) -> None:
        args = meta_data["args"]
        url = trajectory.states[-1]["info"]["page"].url
        state_idx = len(trajectory.states) - 1
        task_id = meta_data["task_id"]

        # Get hard evaluation for current trajectory == score given by evaluator functions from benchmark
        score = hard_verify(trajectory, meta_data, intent)

        # Create entry for current task if not present
        if task_id not in self.score_per_round:
            self.score_per_round[task_id] = {}
            self.score_per_round[task_id]["intent"] = intent
            self.score_per_round[task_id]["rounds_per_state"] = {}
            self.score_per_round[task_id]["scores"] = []

        # Update round counts for the current state of current task
        if state_idx not in self.score_per_round[task_id]["rounds_per_state"]:
            self.score_per_round[task_id]["rounds_per_state"][state_idx] = 0
        else:
            self.score_per_round[task_id]["rounds_per_state"][state_idx] += 1

        last_action = trajectory.actions[-1]
        data = {
            "state_idx": state_idx,
            "score": score,
            "url": url,
            "raw_prediction": last_action["raw_prediction"],
            "round": self.score_per_round[task_id]["rounds_per_state"][state_idx],
        }
        if "retrieved_knowledge" in trajectory.states[-1]:
            data.update({"retrieved_knowledge": trajectory.states[-1]["retrieved_knowledge"]})

        if "critique_executor_loop_utterances" in last_action:
            data.update({"critique_raw_response": last_action["critique_executor_loop_utterances"]["critique"][-1]})

        self.score_per_round[task_id]["scores"].append(data)

        # Dump the score per round to a json file
        update_json(file_path=f"{args.result_dir}/{self.json_file}", data=self.score_per_round)
        update_json(file_path=f"{args.result_dir}/{self.json_file}", data=self.score_per_round)
        update_json(file_path=f"{args.result_dir}/{self.json_file}", data=self.score_per_round)
