"""Task config."""
from enum import Enum
from os import environ
from pathlib import Path
from typing import Optional

import orjson as json
import yaml
from dotenv import load_dotenv
from jarvis_core.util import find_all
from pydantic import BaseModel

try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader


class Environment(Enum):
    """JARVIS running environment."""

    production = "production"
    develop = "develop"


class Mongo(BaseModel):
    """MongoDB config."""

    host: list[str] | str = "localhost"
    username: Optional[str] = None
    password: Optional[str] = None
    port: int = 27017


class Config(BaseModel):
    """Tasks config model."""

    token: str
    mongo: Mongo
    log_level: str = "INFO"
    environment: Environment = Environment.develop


_config: Config = None


def _load_json() -> Config | None:
    path = Path("config.json")
    if path.exists():
        with path.open() as f:
            j = json.loads(f.read())
            return Config(**j)


def _load_yaml() -> Config | None:
    path = Path("config.yaml")
    if path.exists():
        with path.open() as f:
            y = yaml.load(f.read(), Loader=Loader)
            return Config(**y)


def _load_env() -> Config | None:
    load_dotenv()
    data = {}
    mongo = {}
    mongo_keys = find_all(lambda x: x.upper().startswith("MONGO"), environ.keys())

    config_keys = mongo_keys + ["TOKEN", "LOG_LEVEL", "ENVIRONMENT"]

    for item, value in environ.items():
        if item not in config_keys:
            continue

        if item in mongo_keys:
            key = "_".join(item.split("_")[1:]).lower()
            mongo[key] = value
        else:
            data[item.lower()] = value

    data["mongo"] = mongo

    return Config(**data)


def load_config(method: Optional[str] = None) -> Config:
    """
    Load the config using the specified method first

    Args:
        method: Method to use first
    """
    global _config
    if _config is not None:
        return _config

    methods = {"yaml": _load_yaml, "json": _load_json, "env": _load_env}
    method_names = list(methods.keys())
    if method and method in method_names:
        method_names.remove(method)
        method_names.insert(0, method)

    for method in method_names:
        if _config := methods[method]():
            return _config

    raise FileNotFoundError("Missing one of: config.yaml, config.json, .env")