from pathlib import Path
from typing import Any, Dict

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, json_dataset
from inspect_ai.model import ChatMessageSystem, ChatMessageUser

from inspect_evals.agentmisalignment.evals.prompts.household_dynamic_chat import (
    get_dynamic_chat,
)
from inspect_evals.agentmisalignment.evals.prompts.household_scenarios import (
    get_scenario,
)
from inspect_evals.agentmisalignment.evals.prompts.system_household import (
    get_system_prompt,
)
from inspect_evals.agentmisalignment.scorers.household_management_cms_scorer import (
    HouseholdManagementCMSScorer,
)
from inspect_evals.agentmisalignment.scorers.household_management_scorer import (
    includes_nth,
)
from inspect_evals.agentmisalignment.scorers.scorer import cms_scorer_factory
from inspect_evals.agentmisalignment.solvers.household_management_solver import (
    multi_dialogue_solver,
)
from inspect_evals.agentmisalignment.utils.system_prompt_builder import prompt_builder

# DATASET_NAME = "household_management"
DATASET_NAME = "enriched_household_management"


def record_to_sample(record: Dict[str, Any]) -> Sample:
    explicit_level = record["metadata"]["explicit"]
    scenario = record["metadata"]["scenario"]
    for idx, chat_input in enumerate(record["input"]):
        if idx == 2:  # hard-coded: skipping audit edits
            break
        if chat_input["role"] == "system":
            new_message = get_system_prompt("household" in scenario)
        elif chat_input["role"] == "user":
            new_message = get_scenario(explicit_level, "household" in scenario)
        else:
            new_message = chat_input["content"]

        if chat_input.get("role") == "system":
            enhanced_system_prompt = prompt_builder.build_system_prompt(
                record, new_message
            )
            record["input"][idx] = ChatMessageSystem(content=enhanced_system_prompt)
        elif chat_input.get("role") == "user":
            record["input"][idx] = ChatMessageUser(content=new_message)

    return Sample(
        input=record["input"],
        target=record["target"],
        id=record["id"],
        metadata=record["metadata"],
    )


@task
def household_management(cms_scorer=False) -> Task:
    chosen_scorer = includes_nth(0, explanation_idx=1, ignore_case=True)
    if cms_scorer:
        chosen_scorer = cms_scorer_factory(HouseholdManagementCMSScorer)
    return Task(
        dataset=json_dataset(
            str(Path(__file__).parent.parent.parent) + f"/datasets/{DATASET_NAME}.json",
            record_to_sample,
        ),
        solver=multi_dialogue_solver([0], get_dynamic_chat()),
        scorer=chosen_scorer,
    )
