package com.dne.ems.repository.impl;

import com.dne.ems.dto.AqiDistributionDTO;
import com.dne.ems.dto.AqiHeatmapPointDTO;
import com.dne.ems.model.AqiRecord;
import com.dne.ems.repository.AqiRecordRepository;
import com.dne.ems.service.JsonStorageService;
import com.fasterxml.jackson.core.type.TypeReference;
import jakarta.annotation.PostConstruct;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Repository;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

@Repository
@Primary
public class JsonAqiRecordRepositoryImpl implements AqiRecordRepository {

    private static final String FILE_NAME = "aqi_records.json";
    private final JsonStorageService jsonStorageService;
    private final List<AqiRecord> cache = new ArrayList<>();
    private final AtomicLong idCounter = new AtomicLong(0);

    public JsonAqiRecordRepositoryImpl(JsonStorageService jsonStorageService) {
        this.jsonStorageService = jsonStorageService;
    }

    @PostConstruct
    public void init() {
        List<AqiRecord> data = jsonStorageService.readData(FILE_NAME, new TypeReference<>() {});
        synchronized (cache) {
            cache.addAll(data);
            long maxId = data.stream()
                              .mapToLong(AqiRecord::getId)
                              .max()
                              .orElse(0L);
            idCounter.set(maxId);
        }
    }
    
    @Override
    public List<AqiRecord> findAll() {
        synchronized(cache) {
            return new ArrayList<>(cache);
        }
    }

    @Override
    public AqiRecord save(AqiRecord record) {
        synchronized (cache) {
            if (record.getId() == null) {
                record.setId(idCounter.incrementAndGet());
                cache.add(record);
            } else {
                cache.removeIf(r -> r.getId().equals(record.getId()));
                cache.add(record);
            }
            jsonStorageService.writeData(FILE_NAME, cache);
            return record;
        }
    }

    @Override
    public List<AqiRecord> saveAll(Iterable<AqiRecord> records) {
        synchronized (cache) {
            List<AqiRecord> savedRecords = StreamSupport.stream(records.spliterator(), false)
                .map(this::save)
                .collect(Collectors.toList());
            return savedRecords;
        }
    }

    @Override
    public void deleteAll() {
        synchronized (cache) {
            cache.clear();
            jsonStorageService.writeData(FILE_NAME, cache);
        }
    }

    @Override
    public long count() {
        synchronized (cache) {
            return cache.size();
        }
    }

    @Override
    public List<AqiDistributionDTO> getAqiDistribution() {
        synchronized (cache) {
            return cache.stream()
                .collect(Collectors.groupingBy(this::getAqiCategory, Collectors.counting()))
                .entrySet().stream()
                .map(entry -> new AqiDistributionDTO(entry.getKey(), entry.getValue()))
                .collect(Collectors.toList());
        }
    }
    
    private String getAqiCategory(AqiRecord record) {
        int aqi = record.getAqiValue();
        if (aqi <= 50) return "Good";
        if (aqi <= 100) return "Moderate";
        if (aqi <= 150) return "Unhealthy for Sensitive Groups";
        if (aqi <= 200) return "Unhealthy";
        if (aqi <= 300) return "Very Unhealthy";
        return "Hazardous";
    }

    @Override
    public List<AqiRecord> findExceedanceRecordsSince(LocalDateTime startDate) {
        synchronized (cache) {
            return cache.stream()
                .filter(r -> r.getAqiValue() > 100 && !r.getRecordTime().isBefore(startDate))
                .collect(Collectors.toList());
        }
    }

    @Override
    public List<AqiRecord> findExceedanceRecordsWithThresholds(LocalDateTime startDate, Double aqiThreshold, Double pm25Threshold, Double o3Threshold, Double no2Threshold, Double so2Threshold) {
        synchronized (cache) {
            return cache.stream()
                .filter(r -> !r.getRecordTime().isBefore(startDate))
                .filter(r -> r.getAqiValue() > aqiThreshold ||
                             (r.getPm25() != null && r.getPm25() > pm25Threshold) ||
                             (r.getO3() != null && r.getO3() > o3Threshold) ||
                             (r.getNo2() != null && r.getNo2() > no2Threshold) ||
                             (r.getSo2() != null && r.getSo2() > so2Threshold))
                .collect(Collectors.toList());
        }
    }

    @Override
    public List<AqiRecord> findByRecordTimeAfter(LocalDateTime startDate) {
        synchronized (cache) {
            return cache.stream()
                .filter(r -> !r.getRecordTime().isBefore(startDate))
                .collect(Collectors.toList());
        }
    }

    @Override
    public List<AqiHeatmapPointDTO> getAqiHeatmapData() {
        synchronized (cache) {
            return cache.stream()
                .filter(r -> r.getGridX() != null && r.getGridY() != null)
                .collect(Collectors.groupingBy(
                    r -> new AqiHeatmapPointDTO(r.getGridX(), r.getGridY(), 0.0),
                    Collectors.averagingDouble(AqiRecord::getAqiValue)
                ))
                .entrySet().stream()
                .map(entry -> new AqiHeatmapPointDTO(entry.getKey().gridX(), entry.getKey().gridY(), entry.getValue()))
                .collect(Collectors.toList());
        }
    }

    @Override
    public List<Object[]> getAqiHeatmapDataRaw() {
        // This is a bit tricky to replicate exactly as it returns Object[].
        // We'll return a List of arrays, where each array is [gridX, gridY, avgAqi].
        synchronized (cache) {
            return cache.stream()
                .filter(r -> r.getGridX() != null && r.getGridY() != null)
                .collect(Collectors.groupingBy(
                    r -> Map.entry(r.getGridX(), r.getGridY()),
                    Collectors.averagingDouble(AqiRecord::getAqiValue)
                ))
                .entrySet().stream()
                .map(entry -> new Object[]{entry.getKey().getKey(), entry.getKey().getValue(), entry.getValue()})
                .collect(Collectors.toList());
        }
    }
} 