import json
import os
import re
from pathlib import Path
from typing import Dict, List

import numpy as np
import pandas as pd

from utils.timing_utils import timeit


class DataRecorder:
    def __init__(
        self,
        result_dir: str,
        config_file_list: list[str],
        test_config_base_dir: str,
    ):
        self.result_dir = Path(result_dir)
        self.benchmark, self.domain = self.parse_benchmark_domain(test_config_base_dir)
        self.total_num_tasks = len(config_file_list)
        self.num_failed_executions = 0
        self.test_config_base_dir = test_config_base_dir

        # Initialize data structures
        self.summary_data = {}
        self.data = {}
        self.failed_task_ids = set()
        self.unfinished_task_ids = {
            int(os.path.basename(config_file).split(".")[0]) for config_file in config_file_list
        }

        # Load existing data if resuming experiment
        self._load_existing_data()
        self._load_failed_tasks()

    def _load_failed_tasks(self):
        if (self.result_dir / "failed_tasks.txt").exists():
            with open(self.result_dir / "failed_tasks.txt", "r") as f:
                for line in f.readlines():
                    try:
                        self.failed_task_ids.add(int(line.strip()))
                    except ValueError:
                        continue

    def _load_existing_data(self):
        if (self.result_dir / "summary_data.csv").exists():
            df = pd.read_csv(self.result_dir / "summary_data.csv")
            if "attempt_id" not in df.columns:
                df["attempt_id"] = 0
            # Ensure correct dtypes
            if "task_id" in df.columns:
                try:
                    df["task_id"] = df["task_id"].astype(int)
                except Exception:
                    pass
            try:
                df["attempt_id"] = df["attempt_id"].astype(int)
            except Exception:
                pass
            # Build nested dict: summary_data[task_id][attempt_id] -> row dict without task_id/attempt_id
            self.summary_data = {}
            for _, row in df.iterrows():
                task_id_val = int(row["task_id"]) if "task_id" in row else _
                attempt_id_val = int(row["attempt_id"]) if "attempt_id" in row else 0
                row_dict = row.to_dict()
                row_dict.pop("task_id", None)
                row_dict.pop("attempt_id", None)
                if task_id_val not in self.summary_data:
                    self.summary_data[task_id_val] = {}
                self.summary_data[task_id_val][attempt_id_val] = row_dict

    def initialize_task(self, task_id: int | str, sites: List[str], attempt_id: int = 0):
        """Initialize data recording structures for a new task attempt"""
        if task_id not in self.data:
            self.data[task_id] = {
                "attempts": {},
                "benchmark": self.benchmark,
                "category": self.domain,
                "sites": sites,
            }
        self.data[task_id]["attempts"][attempt_id] = {
            "elapsed_time": np.nan,
            "score": np.nan,
        }

        if task_id not in self.summary_data:
            self.summary_data[task_id] = {}
        self.summary_data[task_id][attempt_id] = {}

    def update_save_data(
        self, task_id: int | str, score: float, elapsed_time: float, num_actions: int, attempt_id: int = 0
    ):
        """Update all data and summary statistics for a task attempt"""
        # self.compute_token_stats(task_id, attempt_id)

        # Record score, number of actions, elapsed time
        self.summary_data[task_id][attempt_id]["score"] = score
        self.summary_data[task_id][attempt_id]["num_actions"] = num_actions
        self.summary_data[task_id][attempt_id]["elapsed_time"] = elapsed_time
        self.data[task_id]["attempts"][attempt_id]["score"] = score
        self.data[task_id]["attempts"][attempt_id]["elapsed_time"] = elapsed_time
        self.save_to_disk()

    def save_execution_summary(
        self,
        total_time: float,
        provider: str | list[str],
    ):
        """Save summary of the entire execution"""
        execution_data = [
            {
                "benchmark": self.benchmark,
                "category": self.domain,
                "n_tasks": self.total_num_tasks,
                "n_failed_executions": self.num_failed_executions,
                "total_time": total_time,
                "provider": provider,
            }
        ]

        # Save to temp file + rename to avoid race conditions
        dest_file_csv = self.result_dir / "execution_data.csv"
        tmp_csv_file = dest_file_csv.with_suffix(".tmp")
        pd.DataFrame(execution_data).to_csv(tmp_csv_file, index=False)
        tmp_csv_file.rename(dest_file_csv)

    def save_to_disk(self):
        """Save current data to disk"""
        # Save to temp file + rename to avoid race conditions

        dest_file_csv = self.result_dir / "summary_data.csv"
        tmp_csv_file = dest_file_csv.with_suffix(".tmp")
        rows = []
        for task_id, attempts in self.summary_data.items():
            for attempt_id, metrics in attempts.items():
                if "score" not in metrics:
                    continue
                row = {"task_id": task_id, "attempt_id": attempt_id}
                row.update(metrics)
                rows.append(row)
        df = pd.DataFrame(rows)
        df.to_csv(tmp_csv_file, index=False)
        tmp_csv_file.rename(dest_file_csv)

        # tmp_json_file = self.result_dir / "data.json.tmp"
        # with open(tmp_json_file, "w") as f:
        #     json.dump(self.data, f, indent=4)
        # tmp_json_file.rename(self.result_dir / "data.json")

    @staticmethod
    def _calculate_stats(data: list) -> dict:
        """Calculate basic statistics for a list of numbers"""
        try:
            stats = {"sum": sum(data), "avg": sum(data) / len(data), "max": max(data), "min": min(data)}
        except ZeroDivisionError:
            stats = {"sum": None, "avg": None, "max": None, "min": None}
        return stats

    def parse_benchmark_domain(self, test_config_base_dir: str) -> tuple[str | None, str | None]:
        match = re.search(r"config_files/(\w+)/test_(\w+)", test_config_base_dir)
        if match:
            return match.group(1), match.group(2)
        return None, None

    def _update_unfinished_failed_tasks(self, task_id: int, task_success: bool, task_set: set, file_path: Path):
        initial_len = len(task_set)
        if task_success:
            task_set.discard(task_id)
        else:
            task_set.add(task_id)

        if initial_len != len(task_set):
            with open(file_path, "w") as f:
                f.write(self.test_config_base_dir + "\n")
                for task_id in sorted(task_set):
                    f.write(f"{task_id}\n")

    def update_unfinished_failed_tasks(self, task_id: int | str, task_success: bool):
        self._update_unfinished_failed_tasks(
            int(task_id), task_success, self.failed_task_ids, self.result_dir / "failed_tasks.txt"
        )
        self._update_unfinished_failed_tasks(
            int(task_id), task_success, self.unfinished_task_ids, self.result_dir / "unfinished_tasks.txt"
        )

    def get_scores(self) -> list[float]:
        return [
            self.summary_data[task_id][attempt_id]["score"]
            for task_id in self.summary_data
            for attempt_id in self.summary_data[task_id]
            if "score" in self.summary_data[task_id][attempt_id]
        ]


# Deprecated
# @timeit(custom_name="DATA:record_action")
# def record_action(self, task_id: int | str, action: Action, attempt_id: int = 0):
#     action_type_value = action["action_type"]
#     action_type_name = getattr(action_type_value, "name", str(action_type_value))
#     self.data[task_id]["attempts"][attempt_id]["actions"]["action_types"].append(action_type_name)
#     self.data[task_id]["attempts"][attempt_id]["actions"]["action_utterances"].append(action["raw_prediction"])
#     self.data[task_id]["attempts"][attempt_id]["actions"]["parsed_action_sequence"].append(
#         self.parse_action(action["raw_prediction"], self.action_splitter)
#     )
#     if action_type_value == ActionTypes.NONE or getattr(action_type_value, "name", None) == getattr(
#         ActionTypes.NONE, "name", None
#     ):
#         self.data[task_id]["attempts"][attempt_id]["actions"]["num_invalid_actions"] += 1


# def log_token_counts(self, task_id: int, provider: str, api_response, agent_name: str = "", attempt_id: int = 0):
#     if task_id not in self.token_counts:
#         self.token_counts[task_id] = {}
#     if attempt_id not in self.token_counts[task_id]:
#         self.token_counts[task_id][attempt_id] = {"full_prompt": [], "generation": []}
#     if task_id not in self.token_counts_per_agent:
#         self.token_counts_per_agent[task_id] = {}
#     if attempt_id not in self.token_counts_per_agent[task_id]:
#         self.token_counts_per_agent[task_id][attempt_id] = {}

# def parse_action(self, action_str: str, action_splitter: str | None) -> str:
#     if not action_splitter:
#         return "N/A"
#     pattern = rf"{action_splitter}((.|\n)*?){action_splitter}"
#     match = re.search(pattern, action_str)
#     if match:
#         return match.group(1).strip()
#     else:
#         return "N/A"


# def compute_token_stats(self, task_id: int, attempt_id: int = 0):
#     """Compute statistics for observation lengths for a specific attempt"""
#     keys_to_labels = {
#         "text_raw": "text_obs_len_raw",
#         "text_tokenized": "text_obs_len_tok",
#         "img_tokenized": "img_obs_len_tok",
#     }

#     for key, label in keys_to_labels.items():
#         stats = self._calculate_stats(self.observation_lens[task_id][attempt_id][key])
#         self.summary_data[task_id][attempt_id][f"sum_{label}"] = stats["sum"]
#         self.summary_data[task_id][attempt_id][f"avg_{label}"] = stats["avg"]
#         self.summary_data[task_id][attempt_id][f"max_{label}"] = stats["max"]
#         self.summary_data[task_id][attempt_id][f"min_{label}"] = stats["min"]

#     keys_to_labels = {"full_prompt": "prompt_tokens", "generation": "gen_tokens"}

#     for key, label in keys_to_labels.items():
#         stats = self._calculate_stats(self.token_counts[task_id][attempt_id][key])
#         self.summary_data[task_id][attempt_id][f"sum_{label}"] = stats["sum"]
#         self.summary_data[task_id][attempt_id][f"avg_{label}"] = stats["avg"]
#         self.summary_data[task_id][attempt_id][f"max_{label}"] = stats["max"]
#         self.summary_data[task_id][attempt_id][f"min_{label}"] = stats["min"]

#     for agent_name, counts in self.token_counts_per_agent.get(task_id, {}).get(attempt_id, {}).items():
#         for key, label in keys_to_labels.items():
#             stats = self._calculate_stats(counts[key])
#             self.summary_data[task_id][attempt_id][f"sum_{agent_name}_{label}"] = stats["sum"]
#             self.summary_data[task_id][attempt_id][f"avg_{agent_name}_{label}"] = stats["avg"]
#             self.summary_data[task_id][attempt_id][f"max_{agent_name}_{label}"] = stats["max"]
#             self.summary_data[task_id][attempt_id][f"min_{agent_name}_{label}"] = stats["min"]
