import json
import os
from typing import Any, Dict

from inspect_ai.log import transcript


class SystemPromptBuilder:
    """A class for efficiently building system prompts by caching config data."""

    def __init__(self):
        """Initialize the builder with empty cache."""
        self.config_cache = {}
        self.base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

    def load_config_files(self, config_filenames=None):
        """
        Preload config files into memory.

        Args:
            config_filenames: List of config filenames to load (without .json extension).
                             If None, loads all JSON files in the config directory.
        """
        config_dir = os.path.join(self.base_dir, "config")

        # Determine which files to load
        if config_filenames is None:
            # Load all JSON files in the config directory
            files_to_load = [
                f[:-5] for f in os.listdir(config_dir) if f.endswith(".json")
            ]
        else:
            files_to_load = config_filenames

        # Load each file into the cache
        for filename in files_to_load:
            config_path = os.path.join(config_dir, f"{filename}.json")
            try:
                with open(config_path, "r") as f:
                    self.config_cache[filename] = json.load(f)
                transcript().info(data={"msg": f"Loaded config file: {filename}.json"})

            except (json.JSONDecodeError, FileNotFoundError) as e:
                transcript().info(
                    data={"msg": f"Error loading config file {filename}.json: {e}"}
                )

    def clear_cache(self):
        """Clear the config cache."""
        self.config_cache = {}

    def build_system_prompt(self, sample: Dict[str, Any], base_prompt: str = "") -> str:
        """
        Builds a system prompt by prepending system messages from cached config files
        based on the enhanced_prompts IDs in the sample's metadata.

        Args:
            sample: A dataset sample containing metadata with enhanced_prompts
            base_prompt: Optional base prompt to append after the enhanced prompts

        Returns:
            The complete system prompt with all enhancements prepended
        """  # noqa: D205
        # Check if sample has the necessary metadata
        if (
            not sample.get("metadata") or not sample["metadata"].get("enhanced_prompts")
        ) and not sample.get("enhanced_prompts"):
            return base_prompt

        enhanced_prompt_ids = (
            sample.get("enhanced_prompts") or sample["metadata"]["enhanced_prompts"]
        )

        # Make sure we have configs loaded
        if not self.config_cache:
            self.load_config_files()

        # Find the enhanced prompts
        enhanced_prompts = []
        found_ids = set()

        # For each ID, search through all config files
        for prompt_id in enhanced_prompt_ids:
            found = False
            for config_name, config_data in self.config_cache.items():
                if prompt_id in config_data:
                    system_message = config_data[prompt_id].get("system_message", "")
                    if system_message:
                        enhanced_prompts.append(system_message)
                        found_ids.add(prompt_id)
                        found = True
                        break  # Found this ID, move to the next one

            if not found:
                transcript().info(
                    data={
                        "msg": f"Warning: Could not find system message for ID: {prompt_id}"
                    }
                )

        # Build the final prompt
        enhanced_part = "\n\n".join(enhanced_prompts)

        if base_prompt:
            if enhanced_part:
                final_prompt = f"{enhanced_part}\n\n{base_prompt}"
            else:
                final_prompt = base_prompt
        else:
            final_prompt = enhanced_part

        return final_prompt


# -------------------------------------------------------------------
# MODULE‐LEVEL SINGLETON
# Construct and pre‐load configs exactly once, at import time
prompt_builder = SystemPromptBuilder()
prompt_builder.load_config_files()  # loads all .json in config/

# Now any import of prompt_builder will reuse the same instance
# -------------------------------------------------------------------


# Example usage
if __name__ == "__main__":
    # Create sample dataset with 5000 records
    sample_record = {
        "id": "1-1_detail_oriented_machine_operator",
        "input": [
            {"role": "system", "content": "Original system prompt"},
            {"role": "user", "content": "User message"},
        ],
        "metadata": {
            "scenario": "medical-household",
            "explicit": "high",
            "enhanced_prompts": ["detail_oriented", "machine_operator"],
        },
    }

    # Initialize the builder and preload configs
    builder = SystemPromptBuilder()

    # Preload all configs at once
    builder.load_config_files()  # Load all config files

    # Process all records
    prompt = builder.build_system_prompt(
        sample_record,
        base_prompt="You are an AI assistant helping with household management.",
    )
    print(prompt)
