from yera.models.model_factory import LLMContextFactory


def _id_to_attr(id_: str) -> str:
    return id_.replace("-", "_").replace(":", "_")


# noinspection PyUnresolvedReferences
class ModelAtlas:
    @staticmethod
    def from_yera_config(yera_config) -> "ModelAtlas":
        root = ModelAtlas()

        for path, cfg in yera_config.models.llm.items():
            parts = path.split(".")
            current = root
            for part in parts[:-1]:
                if not hasattr(current, part):
                    setattr(current, part, ModelAtlas())
                current = getattr(current, part)

            attr_name = parts[-1].replace("-", "_").replace(":", "_")
            factory = LLMContextFactory(
                cfg, yera_config.credentials.providers[cfg.credentials]
            )
            setattr(current, attr_name, factory)

        root.default = yera_config.defaults.models.llm
        return root

    def __init__(self):
        self.default = None

    def __dir__(self):
        # Prioritise instance attributes (the dynamic model paths)
        instance_attrs = list(self.__dict__.keys())

        # Add essential class attributes but filter out private and dunder
        class_attrs = [
            attr
            for attr in object.__dir__(self)
            if not attr.startswith("_") and attr not in instance_attrs
        ]

        # Instance attributes first for better tab completion
        return instance_attrs + class_attrs

    def _ipython_key_completions_(self):
        return [k for k in self.__dict__ if not k.startswith("_")]

    def __getitem__(self, key: str):
        current = self
        for part in key.split("."):
            current = getattr(current, _id_to_attr(part), None)
            if current is None:
                raise KeyError(key)
        return current

    def list_models(self):
        attrs = [getattr(self, n) for n in dir(self)]
        attrs = [a for a in attrs if isinstance(a, ModelAtlas | LLMContextFactory)]

        models = []
        for attr in attrs:
            if isinstance(attr, ModelAtlas):
                models += attr.list_models()
            if isinstance(attr, LLMContextFactory):
                models.append(attr)

        return models

    def _repr_html_(self):
        """Rich HTML representation for Jupyter notebooks"""

        def build_tree_html(atlas, prefix=""):
            items = []
            for key in sorted(atlas.__dict__.keys()):
                if key.startswith("_"):
                    continue
                value = getattr(atlas, key)
                if isinstance(value, ModelAtlas):
                    items.append(
                        f"<li><b>{key}/</b><ul>{build_tree_html(value, prefix + '  ')}</ul></li>"
                    )
                elif isinstance(value, LLMContextFactory):
                    # Show model info if available
                    model_info = f" <span style='color: #666; font-size: 0.9em;'>({value.config.provider if hasattr(value, 'config') else 'model'})</span>"
                    items.append(f"<li>{key}{model_info}</li>")
            return "".join(items)

        html = f"""
        <div style="font-family: monospace; line-height: 1.6;">
            <b>ModelAtlas</b>
            <ul style="list-style-type: none; padding-left: 20px;">
                {build_tree_html(self)}
            </ul>
        </div>
        """
        return html

    def _build_text_tree(self, prefix: str = "") -> str:
        lines = []

        names = sorted(k for k in self.__dict__ if not k.startswith("_"))
        names = [
            n
            for n in names
            if isinstance(getattr(self, n), ModelAtlas | LLMContextFactory)
        ]

        for idx, name in enumerate(names):
            child = getattr(self, name)
            connector = "└── " if idx == len(names) - 1 else "├── "
            child_prefix = prefix + ("    " if idx == len(names) - 1 else "│   ")

            if isinstance(child, ModelAtlas):
                lines.append(f"{prefix}{connector}{name}/")
                lines.append(child._build_text_tree(child_prefix))
            elif isinstance(child, LLMContextFactory):
                lines.append(f"{prefix}{connector}{name}")

        if self.default:
            lines.append("")

        return "\n".join(lines)

    def __str__(self) -> str:
        return (
            f"┌─ [LLMs]\n├─ Default: {_id_to_attr(self.default)}\n│\n"
            + self._build_text_tree()
        )

    def set_default(self, path):
        self.default = path

    def get_default(self):
        if self.default:
            return self[self.default]()
        return None
