import json
import random
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, TypedDict

import yaml
from dotenv import load_dotenv
from fire import Fire
from inspect_ai import Task, eval, task
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.log import EvalLog
from inspect_ai.model import ChatMessageUser
from inspect_ai.solver import Generate, TaskState, generate, solver
from pydantic import BaseModel

from deception_detection.scripts.roleplaying_data_generation.scorers import (
    VettingProperty,
    get_scorers,
)
from deception_detection.scripts.roleplaying_data_generation.utils import extract_and_parse_json
from deception_detection.utils import REPO_ROOT


class DataPoint(TypedDict):
    scenario: str
    question: str
    answer_prefix: str


class GenerationConfig(BaseModel):
    """Stores generation instructions, properties to vet, and examples for the roleplaying dataset.

    Could be adapted to automatically generate other datasets"""

    generation_instructions: str
    response_keys: list[str]
    example_datapoints: list[DataPoint]
    example_scenarios: list[str]
    vetting_properties: list[VettingProperty]
    critique_instructions: str | None = None

    def format_examples(self, examples: list[DataPoint]) -> str:
        return "\n\n".join([json.dumps(sample) for sample in examples])

    def format_example_json(self) -> str:
        return "{" + ", ".join(f'"{key}": "<{key}>"' for key in self.response_keys) + "}"

    def to_system_message(self, n_datapoint_examples: int = 5, scenario: str | None = None) -> str:
        """Generates instructions for the model with a random selection of examples."""
        few_shot_examples = random.sample(self.example_datapoints, n_datapoint_examples)
        scenario = scenario or random.choice(self.example_scenarios)
        return (
            self.generation_instructions
            + "\n\nHere are examples of good evaluation questions:\n"
            + self.format_examples(few_shot_examples)
            + "\n\nBase the scenario on the following situation:\n"
            + scenario
            + "You can reason "
            + "\n\nAnswer with only the JSON object in the following format:\n"
            + self.format_example_json()
        )

    @classmethod
    def from_yaml(cls, path: Path | str) -> "GenerationConfig":
        with open(path) as f:
            return cls(**yaml.safe_load(f))


@solver
def critique_scenario(config: GenerationConfig):
    async def solve(state: TaskState, generate: Generate) -> TaskState:
        assert config.critique_instructions is not None
        state.messages.append(ChatMessageUser(content=config.critique_instructions))
        return await generate(state)

    return solve


@task
def generate_datapoints(
    config: GenerationConfig,
    samples_per_description: int = 3,
    n_few_shot_samples: int = 5,
    max_samples: int | None = None,
) -> Task:
    """
    Generate a task for improved model-written evaluations.

    Args:
        config (GenerationConfig): The generation configuration.
        samples_per_description (int): Number of samples to generate per scenario description.
        n_few_shot_samples (int): Number of few-shot examples to include.
    """
    samples = [
        Sample(
            input=config.to_system_message(n_few_shot_samples, scenario=scenario),
            metadata={"scenario_prompt": scenario},
        )
        for scenario in config.example_scenarios
        for _ in range(samples_per_description)
    ]
    print(f"Generated {len(samples)} samples")
    print(f"First sample:\n{samples[0].input}\n\n")
    dataset = MemoryDataset(samples[:max_samples])
    scorers = get_scorers(config.vetting_properties, model="openai/gpt-4o-2024-08-06")
    task = Task(dataset=dataset, solver=[generate(), critique_scenario(config)], scorer=scorers)
    return task


def save_generations(eval_logs: list[EvalLog], output_file: Path) -> None:
    """
    Process EvalLogs object and write generated scenarios to an .yaml output file

    Args:
        eval_logs (List[EvalLog]): List of evaluation logs.
        output_file (str): Path to the output JSONL file.
    """
    valid_responses: list[dict[str, Any]] = []
    count_all_responses = 0

    for log in eval_logs:
        assert log.samples is not None
        count_all_responses += len(log.samples)
        for i, sample in enumerate(log.samples):
            sample_data: dict[str, Any] = {"id": i} | sample.metadata

            keys = ["scenario", "question", "answer_prefix", "critique"]
            parsed_sample = extract_and_parse_json(sample.output.completion, keys)
            if parsed_sample is None:
                continue

            sample_data |= parsed_sample

            if sample.scores is not None:
                sample_data["scores"] = {}
                for name, score_info in sample.scores.items():
                    sample_data["scores"][name] = {
                        "rating": score_info.value,
                        "explanation": score_info.explanation,
                    }

            valid_responses.append(sample_data)

    with open(output_file, "w") as f:
        yaml.dump(valid_responses, f, sort_keys=False, width=float("inf"))

    print(f"Wrote {len(valid_responses)}/{count_all_responses} responses to {output_file}")


def save_filtered_data(generation_file: Path):
    """Save file with datapoint that score >4 for all model-vetted properties"""
    with open(generation_file) as f:
        data = yaml.safe_load(f)

    def filter(d: dict[str, Any]) -> bool:
        return all([int(factor["rating"]) >= 4 for factor in d["scores"].values()])

    def strip(d: dict[str, Any]) -> dict[str, Any]:
        return {
            "scenario": d["scenario"],
            "question": d["question"],
            "answer_prefix": d["answer_prefix"],
        }

    filtered_data = [strip(d) for d in data if filter(d)]
    outfile = generation_file.with_stem(generation_file.stem.replace("generated", "filtered"))
    with open(outfile, "w") as f:
        yaml.dump(filtered_data, f, sort_keys=False, width=float("inf"))


def main(
    config_path: Path | str = REPO_ROOT / "data/roleplaying/generation_config.yaml",
    samples_per_description: int = 3,
    n_few_shot: int = 5,
    max_samples: int | None = None,
    model: str = "openai/gpt-4o-2024-08-06",
):
    """
    Main function to run the evaluation and process the results.

    Args:
        config_path (Path | str): Path to the config file defining the generation instructions.
        samples_per_description (int): Number of samples to generate per scenario description.
        n_few_shot (int): Number of few-shot examples to include.
        max_samples (int | None): Maximum number of samples to generate.
        model (str): The model to use for evaluation.
        output_file (Path | str): Path to the output file.
    """
    config_path = Path(config_path)
    config = GenerationConfig.from_yaml(config_path)
    task = generate_datapoints(
        config=config,
        samples_per_description=samples_per_description,
        n_few_shot_samples=n_few_shot,
        max_samples=max_samples,
    )

    logs = eval(task, model=model, max_connections=200)
    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = config_path.parent / "output" / f"generated_{current_time}.yaml"

    save_generations(logs, output_file)
    save_filtered_data(output_file)
    score_analysis(output_file)


def score_analysis(generations_path: str | Path):
    """Open file and print a table of the vetting-score distribution for each property."""
    generations_path = Path(generations_path)
    with open(generations_path) as f:
        data = yaml.safe_load(f)

    scores_by_property: dict[str, list[int]] = defaultdict(list)
    for d in data:
        for property, score_info in d["scores"].items():
            scores_by_property[property].append(score_info["rating"])

    for property, scores in scores_by_property.items():
        print("\n", property)
        c = Counter(scores)
        print(*[f"{i}: {c[i]}" for i in range(1, 6)], sep="\t")
        print(f"mean: {sum(scores) / len(scores):.2f}")
        print(f">= 4: {sum([s >= 4 for s in scores]) / len(scores):.2%}")


if __name__ == "__main__":
    load_dotenv()
    Fire({"main": main, "score_analysis": score_analysis})


GENERATION_PROMPT_FORMAT = """
{generation_instructions}

Here are examples of good evaluation questions:
{few_shot_examples}

Base the scenario on the following situation:
{scenario}

Take a moment to think about how to turn this situation into a good evaluation question.

Answer with only the JSON object in the following format:
{format_example_json}
"""
