import logging
import os
from typing import Dict, List

import yaml
from pydantic import BaseModel


def load_yaml_file(filename: str) -> dict:
    """
    Load a YAML file from local or script directory and return its contents as a dict.
    """
    local_path = os.path.abspath(filename)
    script_dir = os.path.dirname(os.path.abspath(__file__))
    script_path = os.path.join(script_dir, filename)

    if os.path.exists(local_path):
        logging.info(f"Loaded {filename} from local directory: {local_path}")
        path_to_load = local_path
    elif os.path.exists(script_path):
        logging.info(f"Loaded {filename} from script directory: {script_path}")
        path_to_load = script_path
    else:
        raise FileNotFoundError(
            f"File not found in local or script directory: {filename}"
        )

    with open(path_to_load, "r") as f:
        return yaml.safe_load(f)


class DirectQueryConfig(BaseModel):
    """
    Configuration model for direct-query CLI.
    Contains a list of hint files, a prompt string, and a model name.
    """

    hint_files: list[str]
    prompt: str
    modify_prompt: str
    model_name: str = "gpt-4-1106-preview"
    docker_image: str = ""


def load_config(
    config_path: str = "direct-query-config.yaml",
) -> DirectQueryConfig:
    """
    Load configuration from a YAML file and return a DirectQueryConfig instance.
    Sets default model_name to gpt-4-1106-preview if not present.
    """
    data = load_yaml_file(config_path)
    return DirectQueryConfig(**data)


class PlanQueryConfig(BaseModel):
    """
    Configuration model for direct-query CLI.
    Contains a list of hint files, a prompt string, and a model name.
    """

    hint_files: Dict[str, List[str]]
    prompts: Dict[str, str]
    model_name: str


def load_plan_config(
    config_path: str = "plan-query-config.yaml",
) -> PlanQueryConfig:
    """
    Load configuration from a YAML file and return a DirectQueryConfig instance.
    Sets default model_name to gpt-4-1106-preview if not present.
    """
    data = load_yaml_file(config_path)
    return PlanQueryConfig(**data)
