package com.atlassian.bitbucket.mesh.platform;

import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.springframework.stereotype.Component;

import javax.annotation.Nonnull;
import javax.management.*;
import java.lang.management.ManagementFactory;
import java.util.Optional;
import java.util.OptionalDouble;
import java.util.OptionalLong;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;

import static org.slf4j.LoggerFactory.getLogger;

@Component("operatingSystemService")
public class DefaultOperatingSystemService implements OperatingSystemService {

    private static final Logger log = getLogger(DefaultOperatingSystemService.class);

    private static final String ATTR_SYSTEM_CPU_LOAD = "SystemCpuLoad";
    private static final String ATTR_TOTAL_PHYSICAL_MEMORY_SIZE = "TotalPhysicalMemorySize";

    private final ObjectName osName;
    private final OptionalNumericalAttribute<Double> systemCpuLoad;
    private final OptionalNumericalAttribute<Long> physicalMem;

    public DefaultOperatingSystemService() throws MalformedObjectNameException {
        osName = ObjectName.getInstance("java.lang:type=OperatingSystem");
        systemCpuLoad = new OptionalNumericalAttribute<>(value -> {
            if (value == null) {
                return Optional.empty();
            }

            double asDouble = value.doubleValue();
            if (Double.compare(asDouble, 0D) < 0 || Double.compare(asDouble, 1D) > 0) {
                return Optional.empty();
            }

            return Optional.of(asDouble);
        }, ATTR_SYSTEM_CPU_LOAD);

        physicalMem = new OptionalNumericalAttribute<>(value -> {
            if (value == null) {
                return Optional.empty();
            }

            long asLong = value.longValue();
            if (asLong <= 0L) {
                return Optional.empty();
            }

            return Optional.of(asLong);
        }, ATTR_TOTAL_PHYSICAL_MEMORY_SIZE);

        // on some OSs pre-fetching the CPU load will improve the chances of having a valid value
        // returned when a caller needs it
        OptionalDouble cpuLoad = getCpuLoad();
        log.debug("Pre-fetch of CPU load: {}", cpuLoad.isPresent() ? cpuLoad.getAsDouble() : "not available");
    }

    @Nonnull
    @Override
    public OptionalDouble getCpuLoad() {
        return systemCpuLoad.getValue().map(OptionalDouble::of).orElse(OptionalDouble.empty());
    }

    @Nonnull
    @Override
    public OptionalLong getTotalPhysicalMemory() {
        return physicalMem.getValue().map(OptionalLong::of).orElse(OptionalLong.empty());
    }

    class OptionalNumericalAttribute<V> {

        private final Function<Number, Optional<V>> converter;
        private final AtomicBoolean lastAvailable;
        private final AtomicBoolean loggedFirstAbsentValue;
        private final String name;

        public OptionalNumericalAttribute(Function<Number, Optional<V>> converter, String name) {
            this.converter = converter;
            this.name = name;
            lastAvailable = new AtomicBoolean(true);
            loggedFirstAbsentValue = new AtomicBoolean(false);
        }

        public Optional<V> getValue() {
            try {
                AttributeList list = getAttribute(name);

                if (list.isEmpty()) {
                    if (loggedFirstAbsentValue.compareAndSet(false, true)) {
                        log.debug("Attribute {} on {} is not available", name, osName);
                    }
                    return Optional.empty();
                }

                Attribute att = (Attribute) list.get(0);
                Number value = att.getValue() instanceof Number ? (Number) att.getValue() : null;
                Optional<V> optionalValue = converter.apply(value);

                if (!optionalValue.isPresent()) {
                    if (lastAvailable.compareAndSet(true, false)) {
                        log.debug("Attribute {} on {} is not available. Value returned: {}",
                                name, osName, value);
                    }
                    return Optional.empty();
                }

                if (lastAvailable.compareAndSet(false, true)) {
                    // from my reading some OS JMX attributes may not always be available. It's likely it's only ever
                    // unavailable at the beginning of JVM start-up and subsequently is always available but I'm making
                    // the logging frequency and messaging respect that it may be unavailable at some points
                    log.debug("Attribute {} on {} is once again available. Value returned: {}", name, osName, value);
                }

                return optionalValue;
            } catch (InstanceNotFoundException | ReflectionException e) {
                if (loggedFirstAbsentValue.compareAndSet(false, true)) {
                    log.error("Failed to query attribute {} on {}", name, osName, e);
                }
                return Optional.empty();
            }
        }
    }

    @VisibleForTesting
    AttributeList getAttribute(String name) throws InstanceNotFoundException, ReflectionException {
        return ManagementFactory.getPlatformMBeanServer().getAttributes(
                osName, new String[] {name});
    }
}
