"""
Context and Mode configuration loader
"""

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Self

import yaml
from sensai.util import logging
from sensai.util.string import ToStringMixin

from serena.config.serena_config import SerenaPaths, ToolInclusionDefinition
from serena.constants import (
    DEFAULT_CONTEXT,
    INTERNAL_MODE_YAMLS_DIR,
    SERENA_FILE_ENCODING,
    SERENAS_OWN_CONTEXT_YAMLS_DIR,
    SERENAS_OWN_MODE_YAMLS_DIR,
)

if TYPE_CHECKING:
    pass

log = logging.getLogger(__name__)


def looks_like_yaml_path(s: str) -> bool:
    return os.sep in s or (os.altsep and os.altsep in s) or s.lower().endswith((".yml", ".yaml"))


@dataclass(kw_only=True)
class SerenaAgentMode(ToolInclusionDefinition, ToStringMixin):
    """Represents a mode of operation for the agent, typically read off a YAML file.
    An agent can be in multiple modes simultaneously as long as they are not mutually exclusive.
    The modes can be adjusted after the agent is running, for example for switching from planning to editing.
    """

    name: str
    prompt: str
    """
    a Jinja2 template for the generation of the system prompt.
    It is formatted by the agent (see SerenaAgent._format_prompt()).
    """
    description: str = ""
    _yaml_path: Path | None = field(default=None, repr=False, compare=False)
    """
    Internal field storing the path to the YAML file this mode was loaded from.
    Used to support loading modes from arbitrary file paths.
    """

    def _tostring_includes(self) -> list[str]:
        return ["name"]

    def print_overview(self) -> None:
        """Print an overview of the mode."""
        print(f"{self.name}:\n {self.description}")
        if self.excluded_tools:
            print(" excluded tools:\n    " + ", ".join(sorted(self.excluded_tools)))
        if self.included_optional_tools:
            print(" included optional tools:\n    " + ", ".join(sorted(self.included_optional_tools)))
        if self.fixed_tools:
            print(" fixed tools:\n    " + ", ".join(sorted(self.fixed_tools)))
        if self.prompt:
            print(" defines initial prompt")

    @classmethod
    def from_yaml(cls, yaml_path: str | Path) -> Self:
        """Load a mode from a YAML file."""
        yaml_as_path = Path(yaml_path).resolve()
        with Path(yaml_as_path).open(encoding=SERENA_FILE_ENCODING) as f:
            data = yaml.safe_load(f)
        name = data.pop("name", yaml_as_path.stem)
        return cls(name=name, _yaml_path=yaml_as_path, **data)

    @classmethod
    def get_path(cls, name: str, instance: Self | None = None) -> str:
        """Get the path to the YAML file for a mode.

        :param name: The name of the mode
        :param instance: Optional mode instance. If provided and it has a stored path, that path is returned.
        :return: The path to the mode's YAML file
        """
        # If we have an instance with a stored path, use that
        if instance is not None and instance._yaml_path is not None:
            return str(instance._yaml_path)

        fname = f"{name}.yml"
        custom_mode_path = os.path.join(SerenaPaths().user_modes_dir, fname)
        if os.path.exists(custom_mode_path):
            return custom_mode_path

        own_yaml_path = os.path.join(SERENAS_OWN_MODE_YAMLS_DIR, fname)
        if not os.path.exists(own_yaml_path):
            raise FileNotFoundError(
                f"Mode {name} not found in {SerenaPaths().user_modes_dir} or in {SERENAS_OWN_MODE_YAMLS_DIR}."
                f"Available modes:\n{cls.list_registered_mode_names()}"
            )
        return own_yaml_path

    @classmethod
    def from_name(cls, name: str) -> Self:
        """Load a registered Serena mode."""
        mode_path = cls.get_path(name)
        return cls.from_yaml(mode_path)

    @classmethod
    def from_name_internal(cls, name: str) -> Self:
        """Loads an internal Serena mode"""
        yaml_path = os.path.join(INTERNAL_MODE_YAMLS_DIR, f"{name}.yml")
        if not os.path.exists(yaml_path):
            raise FileNotFoundError(f"Internal mode '{name}' not found in {INTERNAL_MODE_YAMLS_DIR}")
        return cls.from_yaml(yaml_path)

    @classmethod
    def list_registered_mode_names(cls, include_user_modes: bool = True) -> list[str]:
        """Names of all registered modes (from the corresponding YAML files in the serena repo)."""
        modes = [f.stem for f in Path(SERENAS_OWN_MODE_YAMLS_DIR).glob("*.yml") if f.name != "mode.template.yml"]
        if include_user_modes:
            modes += cls.list_custom_mode_names()
        return sorted(set(modes))

    @classmethod
    def list_custom_mode_names(cls) -> list[str]:
        """Names of all custom modes defined by the user."""
        return [f.stem for f in Path(SerenaPaths().user_modes_dir).glob("*.yml")]

    @classmethod
    def load(cls, name_or_path: str | Path) -> Self:
        # If it is a path or looks like a path, load from file
        if isinstance(name_or_path, Path) or looks_like_yaml_path(str(name_or_path)):
            path = Path(name_or_path)
            if path.exists() and path.is_file():
                return cls.from_yaml(name_or_path)
            else:
                raise FileNotFoundError(f"Mode file not found: {path.resolve()}")

        # load from name
        return cls.from_name(str(name_or_path))

    def has_prompt(self) -> bool:
        """
        :return: whether this mode defines a prompt
        """
        return bool(self.prompt and self.prompt.strip())


@dataclass(kw_only=True)
class SerenaAgentContext(ToolInclusionDefinition, ToStringMixin):
    """Represents a context where the agent is operating (an IDE, a chat, etc.), typically read off a YAML file.
    An agent can only be in a single context at a time.
    The contexts cannot be changed after the agent is running.
    """

    name: str
    """the name of the context"""

    prompt: str
    """
    a Jinja2 template for the generation of the system prompt.
    It is formatted by the agent (see SerenaAgent._format_prompt()).
    """

    description: str = ""

    tool_description_overrides: dict[str, str] = field(default_factory=dict)
    """
    maps tool names to custom descriptions, default descriptions are extracted from the tool docstrings.
    """

    _yaml_path: Path | None = field(default=None, repr=False, compare=False)
    """
    Internal field storing the path to the YAML file this context was loaded from.
    Used to support loading contexts from arbitrary file paths.
    """

    single_project: bool = False
    """
    whether to assume that Serena shall only work on a single project in this context (provided that a project is given
    when Serena is started).
    If set to true and a project is provided at startup, the set of tools is limited to those required by the project's
    concrete configuration, and other tools are excluded completely, allowing the set of tools to be minimal.
    The `activate_project` tool will, therefore, be disabled in this case, as project switching is not allowed.
    """

    structured_tool_output: bool | None = None
    """
    whether to use structured output for tools (None = auto)
    """

    def _tostring_includes(self) -> list[str]:
        return ["name"]

    @classmethod
    def from_yaml(cls, yaml_path: str | Path) -> Self:
        """Load a context from a YAML file."""
        yaml_as_path = Path(yaml_path).resolve()
        with yaml_as_path.open(encoding=SERENA_FILE_ENCODING) as f:
            data = yaml.safe_load(f)
        name = data.pop("name", yaml_as_path.stem)
        # Ensure backwards compatibility for tool_description_overrides
        if "tool_description_overrides" not in data:
            data["tool_description_overrides"] = {}
        return cls(name=name, _yaml_path=yaml_as_path, **data)

    @classmethod
    def get_path(cls, name: str, instance: Self | None = None) -> str:
        """Get the path to the YAML file for a context.

        :param name: The name of the context
        :param instance: Optional context instance. If provided and it has a stored path, that path is returned.
        :return: The path to the context's YAML file
        """
        # If we have an instance with a stored path, use that
        if instance is not None and instance._yaml_path is not None:
            return str(instance._yaml_path)

        fname = f"{name}.yml"
        custom_context_path = os.path.join(SerenaPaths().user_contexts_dir, fname)
        if os.path.exists(custom_context_path):
            return custom_context_path

        own_yaml_path = os.path.join(SERENAS_OWN_CONTEXT_YAMLS_DIR, fname)
        if not os.path.exists(own_yaml_path):
            raise FileNotFoundError(
                f"Context {name} not found in {SerenaPaths().user_contexts_dir} or in {SERENAS_OWN_CONTEXT_YAMLS_DIR}."
                f"Available contexts:\n{cls.list_registered_context_names()}"
            )
        return own_yaml_path

    @classmethod
    def from_name(cls, name: str) -> Self:
        """Load a registered Serena context."""
        legacy_name_mapping = {
            "ide-assistant": "claude-code",
        }
        if name in legacy_name_mapping:
            log.warning(
                f"Context name '{name}' is deprecated and has been renamed to '{legacy_name_mapping[name]}'. "
                f"Please update your configuration; refer to the configuration guide for more details: "
                "https://oraios.github.io/serena/02-usage/050_configuration.html#contexts"
            )
            name = legacy_name_mapping[name]
        context_path = cls.get_path(name)
        return cls.from_yaml(context_path)

    @classmethod
    def load(cls, name_or_path: str | Path) -> Self:
        # If it is a path or looks like a path, load from file
        if isinstance(name_or_path, Path) or looks_like_yaml_path(str(name_or_path)):
            path = Path(name_or_path)
            if path.exists() and path.is_file():
                return cls.from_yaml(name_or_path)
            else:
                raise FileNotFoundError(f"Context file not found: {path.resolve()}")

        return cls.from_name(str(name_or_path))

    @classmethod
    def list_registered_context_names(cls, include_user_contexts: bool = True) -> list[str]:
        """Names of all registered contexts (from the corresponding YAML files in the serena repo)."""
        contexts = [f.stem for f in Path(SERENAS_OWN_CONTEXT_YAMLS_DIR).glob("*.yml")]
        if include_user_contexts:
            contexts += cls.list_custom_context_names()
        return sorted(set(contexts))

    @classmethod
    def list_custom_context_names(cls) -> list[str]:
        """Names of all custom contexts defined by the user."""
        return [f.stem for f in Path(SerenaPaths().user_contexts_dir).glob("*.yml")]

    @classmethod
    def load_default(cls) -> Self:
        """Load the default context."""
        return cls.from_name(DEFAULT_CONTEXT)

    def print_overview(self) -> None:
        """Print an overview of the mode."""
        print(f"{self.name}:\n {self.description}")
        if self.excluded_tools:
            print(" excluded tools:\n  " + ", ".join(sorted(self.excluded_tools)))
