from dataclasses import dataclass, field
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",
    "AssistantCompactionSettings",
    "AssistantRunCommandSettings",
    "AssistantSettings",
]


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 AssistantRunCommandDict(TypedDict):
    allowShell: NotRequired[bool]
    shellProgram: NotRequired[str]


class AssistantCompactionDict(TypedDict):
    enabled: NotRequired[bool]
    triggerMessageCount: NotRequired[int]
    preserveRecentMessages: NotRequired[int]
    summaryMaxTokens: NotRequired[int]


class AssistantConfigDict(TypedDict):
    workspaceRoot: NotRequired[str]
    sessionStoragePath: NotRequired[str]
    promptHistoryPath: NotRequired[str]
    allowedRoots: NotRequired[list[str]]
    includeSkillDirectoriesInAllowedRoots: NotRequired[bool]
    maxIterations: NotRequired[int]
    systemPrompt: NotRequired[str]
    runCommand: NotRequired[AssistantRunCommandDict]
    compaction: NotRequired[AssistantCompactionDict]


@dataclass(frozen=True, slots=True)
class AssistantRunCommandSettings:
    allow_shell: bool = False
    shell_program: str = "/bin/bash"


@dataclass(frozen=True, slots=True)
class AssistantCompactionSettings:
    enabled: bool = True
    trigger_message_count: int = 24
    preserve_recent_messages: int = 8
    summary_max_tokens: int = 512


@dataclass(frozen=True, slots=True)
class AssistantSettings:
    workspace_root: Path
    session_storage_path: Path
    prompt_history_path: Path
    allowed_roots: tuple[Path, ...]
    include_skill_directories_in_allowed_roots: bool = True
    max_iterations: int = 100
    system_prompt: str = (
        "You are a thoughtful personal assistant. Help with planning, writing, "
        "summaries, decisions, research guidance, and everyday organization. "
        "Be calm, practical, and concise."
    )
    run_command: AssistantRunCommandSettings = field(
        default_factory=AssistantRunCommandSettings
    )
    compaction: AssistantCompactionSettings = field(
        default_factory=AssistantCompactionSettings
    )


class YIERConfig:
    DEFAULT_DIR = Path.home() / ".yier"
    DEFAULT_ASSISTANT_SYSTEM_PROMPT = (
        "You are a thoughtful personal assistant. Help with planning, writing, "
        "summaries, decisions, research guidance, and everyday organization. "
        "Be calm, practical, and concise."
    )

    @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", {})

    @classmethod
    def resolve_runtime_config_dir(cls, cwd: Path | None = None) -> Path | None:
        """Prefer a project-local .yier.json when present, else fall back to the default dir."""
        base_dir = (cwd or Path.cwd()).resolve()
        local_config = base_dir / ".yier.json"
        if local_config.exists():
            return base_dir
        return None

    @classmethod
    def read_assistant_settings(
        cls,
        cwd: Path | None = None,
        config_dir: Path | None = None,
    ) -> AssistantSettings:
        base_cwd = (cwd or Path.cwd()).resolve()
        payload: AssistantConfigDict = cls.load_config(config_dir).get("assistant", {})

        workspace_root = cls._resolve_assistant_path(
            payload.get("workspaceRoot"),
            base_dir=base_cwd,
            default=base_cwd,
        )
        session_storage_path = cls._resolve_assistant_path(
            payload.get("sessionStoragePath"),
            base_dir=workspace_root,
            default=workspace_root / ".agent_storage" / "sessions",
        )
        prompt_history_path = cls._resolve_assistant_path(
            payload.get("promptHistoryPath"),
            base_dir=workspace_root,
            default=workspace_root / ".agent_storage" / "prompt_history.txt",
        )

        allowed_roots_raw = payload.get("allowedRoots")
        if allowed_roots_raw:
            allowed_roots = cls._normalize_paths(
                cls._resolve_assistant_path(path, base_dir=workspace_root)
                for path in allowed_roots_raw
            )
        else:
            allowed_roots = cls._normalize_paths([workspace_root, session_storage_path])

        run_command_payload = payload.get("runCommand", {})
        compaction_payload = payload.get("compaction", {})

        return AssistantSettings(
            workspace_root=workspace_root,
            session_storage_path=session_storage_path,
            prompt_history_path=prompt_history_path,
            allowed_roots=allowed_roots,
            include_skill_directories_in_allowed_roots=payload.get(
                "includeSkillDirectoriesInAllowedRoots",
                True,
            ),
            max_iterations=payload.get("maxIterations", 100),
            system_prompt=payload.get(
                "systemPrompt",
                cls.DEFAULT_ASSISTANT_SYSTEM_PROMPT,
            ),
            run_command=AssistantRunCommandSettings(
                allow_shell=run_command_payload.get("allowShell", False),
                shell_program=run_command_payload.get("shellProgram", "/bin/bash"),
            ),
            compaction=AssistantCompactionSettings(
                enabled=compaction_payload.get("enabled", True),
                trigger_message_count=compaction_payload.get("triggerMessageCount", 24),
                preserve_recent_messages=compaction_payload.get(
                    "preserveRecentMessages",
                    8,
                ),
                summary_max_tokens=compaction_payload.get("summaryMaxTokens", 512),
            ),
        )

    @staticmethod
    def _resolve_assistant_path(
        raw_path: str | None,
        *,
        base_dir: Path,
        default: Path | None = None,
    ) -> Path:
        if raw_path is None:
            if default is None:
                raise ValueError("A default path is required when raw_path is None")
            return default.resolve()

        candidate = Path(raw_path).expanduser()
        if candidate.is_absolute():
            return candidate.resolve()
        return (base_dir / candidate).resolve()

    @staticmethod
    def _normalize_paths(paths: list[Path] | tuple[Path, ...] | Any) -> tuple[Path, ...]:
        normalized: list[Path] = []
        for path in paths:
            resolved = Path(path).expanduser().resolve()
            if resolved not in normalized:
                normalized.append(resolved)
        return tuple(normalized)
