# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import os
import re
from pathlib import Path
from typing import Annotated, Any, Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, field_validator

from mirage.cache.file.config import CacheConfig, RedisCacheConfig
from mirage.resource.registry import build_resource
from mirage.types import ConsistencyPolicy, MountMode


def _coerce_mount_mode(value):
    if isinstance(value, MountMode):
        return value
    if isinstance(value, str):
        return MountMode(value.lower())
    return value


def _coerce_consistency(value):
    if isinstance(value, ConsistencyPolicy):
        return value
    if isinstance(value, str):
        return ConsistencyPolicy(value.lower())
    return value


_VAR_RE = re.compile(r"\$\{([A-Z_][A-Z0-9_]*)\}")


class _EnvInterpolator:

    def __init__(self, env: dict[str, str], missing: list[str]) -> None:
        self.env = env
        self.missing = missing

    def _sub(self, m: re.Match) -> str:
        name = m.group(1)
        if name not in self.env:
            self.missing.append(name)
            return ""
        return self.env[name]

    def apply(self, value: Any) -> Any:
        if isinstance(value, str):
            return _VAR_RE.sub(self._sub, value)
        if isinstance(value, dict):
            return {k: self.apply(v) for k, v in value.items()}
        if isinstance(value, list):
            return [self.apply(v) for v in value]
        return value


def _interpolate_env(value: Any, env: dict[str, str]) -> Any:
    """Replace ``${VAR}`` placeholders with values from ``env``.

    Args:
        value (Any): scalar, dict, or list to walk.
        env (dict[str, str]): environment mapping to read from.

    Returns:
        Any: ``value`` with every ``${VAR}`` placeholder replaced.

    Raises:
        ValueError: any referenced variable is missing from ``env``.
    """
    missing: list[str] = []
    interp = _EnvInterpolator(env, missing)
    out = interp.apply(value)
    if missing:
        unique_missing = sorted(set(missing))
        raise ValueError(f"missing environment variables: {unique_missing}")
    return out


def _deep_merge(base: dict, override: dict) -> dict:
    result = dict(base)
    for k, v in override.items():
        if (k in result and isinstance(result[k], dict)
                and isinstance(v, dict)):
            result[k] = _deep_merge(result[k], v)
        else:
            result[k] = v
    return result


class RamCacheBlock(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal["ram"] = "ram"
    limit: str | int = "512MB"
    max_drain_bytes: int | None = None


class RedisCacheBlock(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal["redis"]
    limit: str | int = "512MB"
    max_drain_bytes: int | None = None
    url: str = "redis://localhost:6379/0"
    key_prefix: str = "mirage:cache:"


CacheBlock = Annotated[
    RamCacheBlock | RedisCacheBlock,
    Field(discriminator="type"),
]


class MountBlock(BaseModel):
    model_config = ConfigDict(extra="forbid")

    resource: str
    mode: MountMode | None = None
    config: dict[str, Any] = Field(default_factory=dict)

    @field_validator("mode", mode="before")
    @classmethod
    def _v_mode(cls, v):
        if v is None:
            return v
        return _coerce_mount_mode(v)


class WorkspaceConfig(BaseModel):
    model_config = ConfigDict(extra="forbid")

    mounts: dict[str, MountBlock]
    mode: MountMode = MountMode.WRITE
    consistency: ConsistencyPolicy = ConsistencyPolicy.LAZY
    default_session_id: str = "default"
    default_agent_id: str = "default"
    fuse: bool = False
    native: bool = False
    history: int | None = 100
    history_path: str | None = None
    cache: CacheBlock | None = None

    @field_validator("mode", mode="before")
    @classmethod
    def _v_mode(cls, v):
        return _coerce_mount_mode(v)

    @field_validator("consistency", mode="before")
    @classmethod
    def _v_cons(cls, v):
        return _coerce_consistency(v)

    def to_workspace_kwargs(self) -> dict[str, Any]:
        """Produce kwargs ready to splat into ``Workspace(**kwargs)``.

        Returns:
            dict[str, Any]: resource instances, cache config, and
                workspace-level settings, in the shape the
                ``Workspace`` constructor expects.
        """
        resources: dict[str, Any] = {}
        for prefix, block in self.mounts.items():
            prov = build_resource(block.resource, block.config)
            mode = block.mode if block.mode is not None else self.mode
            resources[prefix] = (prov, mode)
        kwargs: dict[str, Any] = {
            "resources": resources,
            "mode": self.mode,
            "consistency": self.consistency,
            "session_id": self.default_session_id,
            "agent_id": self.default_agent_id,
            "fuse": self.fuse,
            "native": self.native,
            "history": self.history,
            "history_path": self.history_path,
        }
        if self.cache is not None:
            kwargs["cache"] = _build_cache_config(self.cache)
        return kwargs


def _build_cache_config(block: RamCacheBlock | RedisCacheBlock) -> CacheConfig:
    if isinstance(block, RedisCacheBlock):
        return RedisCacheConfig(
            limit=block.limit,
            max_drain_bytes=block.max_drain_bytes,
            url=block.url,
            key_prefix=block.key_prefix,
        )
    return CacheConfig(
        limit=block.limit,
        max_drain_bytes=block.max_drain_bytes,
    )


def load_config(source: str | Path | dict,
                env: dict[str, str] | None = None) -> WorkspaceConfig:
    """Load a workspace config from a YAML / JSON file or a raw dict.

    Performs ``${VAR}`` env interpolation before validation. If any
    referenced variable is missing, raises with the full list of
    missing names rather than failing lazily on first use.

    Args:
        source (str | Path | dict): path to a YAML / JSON file, or
            an already-parsed dict.
        env (dict[str, str] | None): environment mapping to read for
            interpolation. Defaults to ``os.environ``.

    Returns:
        WorkspaceConfig: validated config object.
    """
    if isinstance(source, (str, Path)):
        text = Path(source).read_text(encoding="utf-8")
        raw = yaml.safe_load(text)
    else:
        raw = dict(source)
    if not isinstance(raw, dict):
        raise ValueError(
            f"config source must be a mapping, got {type(raw).__name__}")
    use_env = env if env is not None else dict(os.environ)
    interpolated = _interpolate_env(raw, use_env)
    return WorkspaceConfig.model_validate(interpolated)


def merge_override(base: WorkspaceConfig,
                   override: str | Path | dict,
                   env: dict[str, str] | None = None) -> WorkspaceConfig:
    """Apply a partial config on top of an existing one.

    Used by ``--clone --override`` and ``--load --override`` to swap
    selected fields (typically resource creds, occasionally a bucket
    or URL) without rewriting the whole config.

    Merge semantics: nested dicts merge by key recursively; leaf
    values replace. Mounts not mentioned in the override survive
    unchanged.

    Args:
        base (WorkspaceConfig): the original config.
        override (str | Path | dict): partial config -- either a path
            to a YAML / JSON file or an already-parsed dict.
        env (dict[str, str] | None): env mapping for ``${VAR}``
            interpolation in the override. Defaults to ``os.environ``.

    Returns:
        WorkspaceConfig: a new validated config with the override
        applied.
    """
    if isinstance(override, (str, Path)):
        text = Path(override).read_text(encoding="utf-8")
        raw = yaml.safe_load(text) or {}
    else:
        raw = dict(override)
    if not isinstance(raw, dict):
        raise ValueError(
            f"override must be a mapping, got {type(raw).__name__}")
    use_env = env if env is not None else dict(os.environ)
    interpolated = _interpolate_env(raw, use_env)
    base_dict = base.model_dump()
    merged = _deep_merge(base_dict, interpolated)
    return WorkspaceConfig.model_validate(merged)
