from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass(frozen=True)
class V2ApiConfig:
    host: str = "0.0.0.0"
    port: int = 8080


@dataclass(frozen=True)
class V2AuthConfig:
    token_env: str = "MVP_INTERNAL_TOKEN"


@dataclass(frozen=True)
class V2SqliteConfig:
    db_path: str


@dataclass(frozen=True)
class V2SchedulerConfig:
    tick_s: int = 5
    retry_interval_s: int = 60
    max_running_tasks: int = 1


@dataclass(frozen=True)
class V2RetentionConfig:
    jobs_trash_after_days: int = 3
    jobs_purge_after_days: int = 7
    janitor_interval_s: int = 3600


@dataclass(frozen=True)
class V2WandbConfig:
    enabled: bool = False
    base_url: str = ""
    api_key_env: str = "WANDB_API_KEY"
    project_suffix: str = "_project"


@dataclass(frozen=True)
class V2TrackingConfig:
    wandb: V2WandbConfig


@dataclass(frozen=True)
class V2SFTPGoConfig:
    enabled: bool = False
    host: str = ""
    sftp_port: int = 2022
    admin_api_base: str = ""
    admin_user: str = "admin"
    admin_password_env: str = "SFTPGO_ADMIN_PASSWORD"


@dataclass(frozen=True)
class V2ServingServeConfig:
    http_port: int = 8000
    proxy_location: str = "HeadOnly"


@dataclass(frozen=True)
class V2ServingLLMConfig:
    accelerator_type: str = ""


@dataclass(frozen=True)
class V2ServingConfig:
    enabled: bool = False
    serve: V2ServingServeConfig = V2ServingServeConfig()
    llm: V2ServingLLMConfig = V2ServingLLMConfig()


@dataclass(frozen=True)
class V2DataConfig:
    user_root: str
    sftpgo: V2SFTPGoConfig
    retention: V2RetentionConfig


@dataclass(frozen=True)
class V2Config:
    api: V2ApiConfig
    auth: V2AuthConfig
    sqlite: V2SqliteConfig
    scheduler: V2SchedulerConfig
    tracking: V2TrackingConfig
    data: V2DataConfig
    serving: V2ServingConfig

    @staticmethod
    def from_root_dict(root: dict[str, Any]) -> "V2Config":
        # New format: root.service
        # Backward compatible: root.v2
        service = root.get("service")
        if service is None:
            service = root.get("v2") or {}
        if not isinstance(service, dict):
            raise ValueError("config.service must be a mapping")

        api = service.get("api") or {}
        auth = service.get("auth") or {}
        sqlite = service.get("sqlite") or {}
        scheduler = service.get("scheduler") or {}

        if not isinstance(api, dict) or not isinstance(auth, dict) or not isinstance(sqlite, dict) or not isinstance(scheduler, dict):
            raise ValueError("config.service.{api,auth,sqlite,scheduler} must be mappings")

        ray = root.get("ray")
        if isinstance(ray, dict) and ray.get("shared_root"):
            shared_root = str(ray.get("shared_root"))
        else:
            shared_root = str(root.get("shared_root") or "/private")

        tracking = root.get("tracking") or {}
        if not isinstance(tracking, dict):
            raise ValueError("config.tracking must be a mapping")
        wandb = tracking.get("wandb") or {}
        if not isinstance(wandb, dict):
            raise ValueError("config.tracking.wandb must be a mapping")

        data = root.get("data") or {}
        if not isinstance(data, dict):
            raise ValueError("config.data must be a mapping")
        sftpgo = data.get("sftpgo") or {}
        retention = data.get("retention") or {}
        if not isinstance(sftpgo, dict) or not isinstance(retention, dict):
            raise ValueError("config.data.{sftpgo,retention} must be mappings")

        serving = root.get("serving") or {}
        if not isinstance(serving, dict):
            raise ValueError("config.serving must be a mapping")
        serving_enabled = bool(serving.get("enabled")) if "enabled" in serving else bool(serving)
        serving_serve = serving.get("serve") or {}
        serving_llm = serving.get("llm") or {}
        if not isinstance(serving_serve, dict) or not isinstance(serving_llm, dict):
            raise ValueError("config.serving.{serve,llm} must be mappings")

        default_db_path = f"{shared_root}/common/db/mvp.sqlite3"
        db_path = str(sqlite.get("db_path") or default_db_path)

        user_root = str(data.get("user_root") or f"{shared_root}/users")

        wandb_enabled = bool(wandb.get("enabled") or False)
        wandb_base_url = str(wandb.get("base_url") or "")
        if wandb_enabled and not wandb_base_url:
            raise ValueError("tracking.wandb.base_url is required when tracking.wandb.enabled=true")

        return V2Config(
            api=V2ApiConfig(
                host=str(api.get("host") or "0.0.0.0"),
                port=int(api.get("port") or 8080),
            ),
            auth=V2AuthConfig(token_env=str(auth.get("token_env") or "MVP_INTERNAL_TOKEN")),
            sqlite=V2SqliteConfig(db_path=db_path),
            scheduler=V2SchedulerConfig(
                tick_s=int(scheduler.get("tick_s") or 5),
                retry_interval_s=int(scheduler.get("retry_interval_s") or 60),
                max_running_tasks=int(scheduler.get("max_running_tasks") or 1),
            ),
            tracking=V2TrackingConfig(
                wandb=V2WandbConfig(
                    enabled=wandb_enabled,
                    base_url=wandb_base_url,
                    api_key_env=str(wandb.get("api_key_env") or "WANDB_API_KEY"),
                    project_suffix=str(wandb.get("project_suffix") or "_project"),
                )
            ),
            data=V2DataConfig(
                user_root=user_root,
                sftpgo=V2SFTPGoConfig(
                    enabled=bool(sftpgo.get("enabled") or False),
                    host=str(sftpgo.get("host") or ""),
                    sftp_port=int(sftpgo.get("sftp_port") or 2022),
                    admin_api_base=str(sftpgo.get("admin_api_base") or ""),
                    admin_user=str(sftpgo.get("admin_user") or "admin"),
                    admin_password_env=str(sftpgo.get("admin_password_env") or "SFTPGO_ADMIN_PASSWORD"),
                ),
                retention=V2RetentionConfig(
                    jobs_trash_after_days=int(retention.get("jobs_trash_after_days") or 3),
                    jobs_purge_after_days=int(retention.get("jobs_purge_after_days") or 7),
                    janitor_interval_s=int(retention.get("janitor_interval_s") or 3600),
                ),
            ),
            serving=V2ServingConfig(
                enabled=serving_enabled,
                serve=V2ServingServeConfig(
                    http_port=int(serving_serve.get("http_port") or 8000),
                    proxy_location=str(serving_serve.get("proxy_location") or "HeadOnly"),
                ),
                llm=V2ServingLLMConfig(
                    accelerator_type=str(serving_llm.get("accelerator_type") or ""),
                ),
            ),
        )
