# Copyright: (c) 2024, Luca Bilke <luca@bil.ke>
# MIT License (see LICENSE)

from __future__ import annotations

import copy
from dataclasses import asdict, dataclass, field, replace
from typing import Any, Callable

import yaml
from ansible.module_utils.basic import AnsibleModule

PROJECTS_DIR = "/var/lib/ez_compose"
BASE_SERVICE_ARGS = {
    "project_name": {
        "type": "str",
        "required": True,
    },
    "name": {
        "type": "str",
        "required": True,
    },
    "image": {
        "type": "str",
    },
    "state": {
        "type": "str",
        "default": "present",
        "choices": ["present", "absent"],
    },
    "defaults": {
        "type": "dict",
    },
}


@dataclass(frozen=True)
class Result:
    changed: bool = False
    diff: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True)
class Settings:
    projects_dir: str = "/usr/local/share/ez_compose/"


@dataclass(frozen=True)
class State:
    module: Any  # Replace Any with the actual type of AnsibleModule if available
    result: Result
    compose_filepath: str
    before: dict[str, Any]
    after: dict[str, Any] = field(default_factory=dict)


def _recursive_update(default: dict[Any, Any], update: dict[Any, Any]) -> dict[Any, Any]:
    for key in update:
        if isinstance(update[key], dict) and isinstance(default.get(key), dict):
            default[key] = _recursive_update(default[key], update[key])

        elif isinstance(update[key], list) and isinstance(default.get(key), list):
            default_set = set(default[key])
            custom_set = set(update[key])
            default[key] = list(default_set.union(custom_set))

        else:
            default[key] = update[key]

    return default


def get_state(module: AnsibleModule) -> State:
    compose_filepath = f"{PROJECTS_DIR}/{module.params['project_name']}/docker-compose.yml"

    try:
        with open(compose_filepath, "r") as fp:
            before = yaml.safe_load(fp)

    except FileNotFoundError:
        before: dict[str, Any] = {}

    return State(
        module=module,
        result=Result(),
        compose_filepath=compose_filepath,
        before=before,
    )


def apply_service_base(state: State) -> State:
    params = state.module.params

    compose = copy.deepcopy(state.before)
    services = compose.get("services", {})

    networks = compose.get("networks", {}).update(
        {
            f"{params['project_name']}_internal": None,
        }
    )

    service = services.get(params["name"], {}).update(
        {
            "container_name": f"{params['project_name']}_{params['name']}",
            "hostname": f"{params['project_name']}_{params['name']}",
            "image": params["image"],
            "restart": "unless-stopped",
            "environment": {},
            "labels": {},
            "volumes": [],
            "networks": {
                f"{params['project_name']}_internal": None,
            },
        }
    )

    services.update(
        {
            "networks": networks,
            params["name"]: service,
        }
    )

    return replace(state, after=compose)


def set_defaults(state: State) -> State:
    params = state.module.params
    compose = copy.deepcopy(state.before)
    services = compose["services"]
    service = services[params["name"]]

    _recursive_update(service, params["defaults"])

    services.update({params["name"]: service})

    return replace(state, after=compose)


def update_service(state: State, update: dict[str, Any]) -> State:
    compose = copy.deepcopy(state.before)
    name = state.module.params["name"]
    _recursive_update(compose["services"][name], update)
    return replace(state, after=compose)


def write_compose(state: State) -> State:
    file = state.compose_filepath

    with open(file, mode="w") as stream:
        yaml.dump(state.after, stream)

    return state


def run_service(
    extra_args: dict[str, Any] = {},
    helper: Callable[[State], State] = lambda x: x,
) -> None:
    module = AnsibleModule(
        argument_spec=BASE_SERVICE_ARGS.update(extra_args),
        supports_check_mode=True,
    )

    state = get_state(module)

    for f in [apply_service_base, set_defaults, helper]:
        state = f(state)

    if module.check_mode:
        module.exit_json(**asdict(state.result))  # type: ignore[reportUnkownMemberType]

    write_compose(state)

    module.exit_json(**asdict(state.result))  # type: ignore[reportUnkownMemberType]