import json
from pathlib import Path
from typing import Any, Literal, TypedDict, NotRequired
from yier_agents.src.log import logger


__all__ = [
    "YIERConfig",
    "McpServerDict",
    "McpServerStatus",
    "StdioMcpServerConfig",
    "HttpMcpServerConfig",
    "SseMcpServerConfig",
]


McpServerStatus = Literal[
    "connected",
    "disabled",
    "failed",
    "needs_auth",
    "needs_client_registration",
]


class StdioMcpServerConfig(TypedDict):
    type: Literal["stdio"]
    command: str
    args: list[str]
    env: NotRequired[dict[str, str]]
    enabled: NotRequired[bool]
    status: NotRequired[McpServerStatus]


class HttpMcpServerConfig(TypedDict):
    type: Literal["http"]
    url: str
    headers: NotRequired[dict[str, str]]
    enabled: NotRequired[bool]
    status: NotRequired[McpServerStatus]


class SseMcpServerConfig(TypedDict):
    type: Literal["sse"]
    url: str
    headers: NotRequired[dict[str, str]]
    enabled: NotRequired[bool]
    status: NotRequired[McpServerStatus]


McpServerDict = StdioMcpServerConfig | HttpMcpServerConfig | SseMcpServerConfig


class YIERConfig:
    DEFAULT_DIR = Path.home() / ".yier"

    @classmethod
    def resolve_config_dir(cls, config_dir: Path | None = None) -> Path:
        if config_dir is None:
            config_dir = cls.DEFAULT_DIR
            config_dir.mkdir(exist_ok=True)
        return config_dir

    @classmethod
    def resolve_config_path(cls, config_dir: Path | None = None) -> Path:
        return cls.resolve_config_dir(config_dir) / ".yier.json"

    @classmethod
    def load_config(cls, config_dir: Path | None = None) -> dict[str, Any]:
        config_dir = cls.resolve_config_dir(config_dir)
        if config_dir != cls.DEFAULT_DIR and not config_dir.exists():
            raise FileNotFoundError(
                f"Yier config directory not found: {config_dir}"
            )
        config_file = cls.resolve_config_path(config_dir)
        try:
            return json.loads(config_file.read_text())
        except FileNotFoundError:
            logger.warning("Yier config file not found.")
            return {}
        except json.JSONDecodeError:
            logger.warning("Yier config file is not valid JSON, skipping.")
            return {}
        except Exception:
            logger.error("Error loading Yier config")
            return {}

    @classmethod
    def read_mcp_config(
        cls, config_dir: Path | None = None
    ) -> dict[str, McpServerDict]:
        return cls.load_config(config_dir).get("mcpServers", {})
