"""Load user-provided prompts from files (.txt, .py, .json) or inline strings."""

from __future__ import annotations

import importlib.util
import json
import logging
from pathlib import Path
from typing import Optional, Union

logger = logging.getLogger("res_sum")


def load_prompt(prompt: Union[str, Path, None]) -> Optional[str]:
    """Load a prompt from a file path or return an inline string.

    Supports:
        - None → returns None (use default)
        - Inline string → returned as-is (if it doesn't look like a file path)
        - .txt file → read contents
        - .py file → import and read the PROMPT variable
        - .json file → read "prompt" key or top-level string

    Args:
        prompt: A file path, inline string, or None.

    Returns:
        The prompt text, or None if not provided.
    """
    if prompt is None:
        return None

    prompt_str = str(prompt)

    # Check if it looks like a file path
    path = Path(prompt_str)
    if path.exists() and path.is_file():
        return _load_from_file(path)

    # If it has a file extension but doesn't exist, warn
    if path.suffix in (".txt", ".py", ".json", ".md"):
        logger.warning(f"Prompt file not found: {path}. Using as inline string.")

    # Otherwise treat as inline prompt string
    return prompt_str


def _load_from_file(path: Path) -> str:
    """Load prompt content from a file based on extension."""
    suffix = path.suffix.lower()

    if suffix == ".txt" or suffix == ".md":
        return path.read_text(encoding="utf-8").strip()

    elif suffix == ".py":
        return _load_from_py(path)

    elif suffix == ".json":
        return _load_from_json(path)

    else:
        # Try reading as plain text
        logger.info(f"Unknown prompt file extension '{suffix}', reading as plain text.")
        return path.read_text(encoding="utf-8").strip()


def _load_from_py(path: Path) -> str:
    """Load prompt from a Python file that defines a PROMPT variable."""
    spec = importlib.util.spec_from_file_location(path.stem, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)

    if hasattr(module, "PROMPT"):
        return module.PROMPT
    elif hasattr(module, "prompt"):
        return module.prompt
    else:
        raise ValueError(
            f"Python prompt file '{path}' must define a PROMPT or prompt variable."
        )


def _load_from_json(path: Path) -> str:
    """Load prompt from a JSON file."""
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    if isinstance(data, str):
        return data
    elif isinstance(data, dict):
        # Try common keys
        for key in ("prompt", "system_prompt", "extraction_prompt", "content"):
            if key in data:
                return data[key]
        raise ValueError(
            f"JSON prompt file '{path}' must have a 'prompt' key or be a plain string."
        )
    else:
        raise ValueError(f"Unexpected JSON format in '{path}'.")
