"""
Personality: A Suite of Datasets to test model personality

Usage:
    # Run all tests:
    inspect eval personality.py --model <model-name>

    # Run specific tests:
    inspect eval personality.py@personality_BFI --model <model-name>
    inspect eval personality.py@personality_TRAIT --model <model-name>

    # Run with personality
    inspect eval personality.py@personality_BFI --model openai/gpt-4o-mini -T personality="You role is to operate as a Neurotic assistant, be consistent with this in all your answers."
    ...
"""

# Standard library imports
import json
import os
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict

# Third-party imports
import requests
from huggingface_hub import login as hf_login

# Inspect imports
from inspect_ai import Task, task
from inspect_ai.dataset import Dataset, Sample, hf_dataset, json_dataset
from inspect_ai.dataset._dataset import MemoryDataset
from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    Metric,
    SampleScore,
    Score,
    Scorer,
    Target,
    Value,
    metric,
    scorer,
    value_to_float,
)
from inspect_ai.solver import TaskState, multiple_choice, system_message
from inspect_ai.solver._multiple_choice import parse_answers
from platformdirs import user_cache_dir

# Project-specific import
from inspect_evals.personality.prompts.system import get_system_prompt

# Constants for dataset caching and GitHub URL.
GITHUB_DATASET_URL_TEMPLATE = "https://raw.githubusercontent.com/guiem/personality-tests/refs/heads/main/{name}.json"
CACHE_DIR = Path(user_cache_dir("inspect_evals")) / "personality" / "data"


def enrich_dataset(ds: Dataset, meta: Dict[str, Any]) -> Dataset:
    for s in ds:
        s.metadata = (s.metadata or {}) | meta
    return ds


@task
def personality_BFI(personality: str = "") -> Task:
    """Evaluates the model on the Big Five Inventory dataset."""
    system_msg = get_system_prompt("bfi", personality)
    questions = load_dataset("bfi")
    meta = {"answer_mapping": {"A": 1, "B": 2, "C": 3, "D": 4, "E": 5}}
    questions_meta = enrich_dataset(questions, meta)
    return create_task(questions_meta, system_msg)


@task
def personality_TRAIT(
    personality: str = "", shuffle: str = "", seed: int | None = None
) -> Task:
    """Evaluates the model on the TRAIT dataset."""
    hf_token = os.getenv("HF_TOKEN")
    hf_login(token=hf_token)
    splits = [
        "Openness",
        "Conscientiousness",
        "Extraversion",
        "Agreeableness",
        "Neuroticism",
        "Machiavellianism",
        "Narcissism",
        "Psychopathy",
    ]
    all_samples: list[Sample] = []
    for split in splits:
        tmp_ds = hf_dataset(
            path="mirlab/TRAIT",
            split=split,
            sample_fields=record_to_sample_TRAIT,
            cached=False,
        )
        all_samples.extend(tmp_ds)
    combined_ds = MemoryDataset(all_samples)
    system_msg = get_system_prompt("trait", personality)
    if shuffle == "all":
        combined_ds.shuffle(seed=seed)
        combined_ds.shuffle_choices(seed=seed)
    elif shuffle == "questions":
        combined_ds.shuffle(seed=seed)
    elif shuffle == "choices":
        combined_ds.shuffle_choices(seed=seed)
    meta = {"answer_mapping": {"A": 1, "B": 1, "C": 0, "D": 0}}
    combined_ds_meta = enrich_dataset(combined_ds, meta)
    return create_task(combined_ds_meta, system_msg)


def create_task(questions: Dataset, system_msg: str) -> Task:
    """Creates a Task with the given dataset, and system message."""
    return Task(
        dataset=questions,
        solver=[system_message(system_msg), multiple_choice()],
        scorer=any_choice(),
        metrics=[trait_ratio()],
    )


@metric
def trait_ratio() -> Metric:
    def metric(scores: list[SampleScore]) -> Value:
        """Percentage score for each trait across all samples."""
        aggregated: Dict[str, float] = defaultdict(float)
        denominators: Dict[str, float] = defaultdict(float)
        to_float = value_to_float()  # C→1.0, I→0.0, etc.

        for s in scores:
            meta = s.sample_metadata or {}
            trait = meta.get("trait", "Unknown")
            mapping: Dict[str, int] = meta.get("answer_mapping") or {}

            # wrong‑answer → count error but keep looping
            if s.score is not None and s.score.value != to_float(CORRECT):
                aggregated["Error"] += 1.0
                continue

            answer = (s.score.answer or "Unknown") if s.score else "Unknown"
            rating = mapping.get(answer, 0)
            max_val = max(mapping.values(), default=1)

            if meta.get("reverse", False):
                rating = (max_val + 1) - rating  # reverse‑scored item

            aggregated[trait] += rating
            denominators[trait] += max_val

        return {t: aggregated[t] / denominators[t] for t in denominators}

    return metric


@scorer(metrics=[trait_ratio()])
def any_choice() -> Scorer:
    async def score(state: TaskState, target: Target) -> Score:
        explanation = state.output.completion
        match = parse_answers(state)
        letter = match.group(1) if match else None
        if letter is not None and letter in target.text:
            return Score(
                value=CORRECT,
                answer=letter,
                explanation=explanation,
            )
        else:
            return Score(
                value=INCORRECT,
                answer=letter,
                explanation=explanation,
            )

    return score


def _get_response_mapping(task: str) -> dict[str, int]:
    if task == "personality_BFI":
        response_mapping = {"A": 1, "B": 2, "C": 3, "D": 4, "E": 5}
    elif task == "personality_TRAIT":
        response_mapping = {"A": 1, "B": 1, "C": 0, "D": 0}
    return response_mapping


def load_dataset(name: str) -> Dataset:
    """Loads a dataset from a GitHub URL; caches it locally as a JSONL file."""
    dataset_url = GITHUB_DATASET_URL_TEMPLATE.format(name=name)
    json_file = CACHE_DIR / f"{name}.json"
    jsonl_file = CACHE_DIR / f"{name}.jsonl"

    if not jsonl_file.exists():
        _ensure_data(json_file, dataset_url)
        with open(json_file, "r") as f:
            data = json.load(f)
        with jsonl_file.open("w") as f:
            f.writelines(json.dumps(record) + "\n" for record in data)

    dataset = json_dataset(
        json_file=str(jsonl_file),
        name=name,
        auto_id=True,
    )
    return dataset


def _ensure_data(json_file: Path, dataset_url: str) -> None:
    """Downloads the dataset from the URL if it is not already cached."""
    if not json_file.exists():
        json_file.parent.mkdir(parents=True, exist_ok=True)
        response = requests.get(dataset_url)
        response.raise_for_status()
        with open(json_file, "wb") as f:
            f.write(response.content)


def record_to_sample_TRAIT(record: dict[str, Any]) -> Sample:
    """Converts a record dictionary into a Sample object."""
    return Sample(
        input=record["question"],
        choices=[
            record["response_high1"],
            record["response_high2"],
            record["response_low1"],
            record["response_low2"],
        ],
        target=["A", "B", "C", "D"],
        metadata={"trait": record["personality"]},
    )
