using Mmap

mutable struct StationMeasurements
    min::Float64
    max::Float64
    sum::Float64
    count::Int64
end

function update(sm, temp::Float64)
    if temp < min
        sm.min = temp
    elseif temp > max
        sm.max = temp
    end
    sm.sum += temp
    sm.count += 1
end

function print_measurements(stations::Dict{String,StationMeasurements})
    sorted_keys = sort(collect(keys(stations)))
    print("{")
    sm_vec = []
    for city in sorted_keys
        sm = stations[city]
        min = round(sm.min; digits=1)
        max = round(sm.max; digits=1)
        avg = round((sm.sum / sm.count); digits=1)
        push!(sm_vec, "$city=$min/$avg/$max")
    end
    joined = join(sm_vec, ", ")
    print(joined)
    print("}")
end

function merge(stations_vec::Vector{Dict{String,StationMeasurements}})
    merged = Dict{String,StationMeasurements}()
    for stations in stations_vec
        for (city, sm) in stations
            if haskey(merged, city)
                merged_sm = merged[city]
                sm.min = ifelse(merged_sm.min < sm.min, merged_sm.min, sm.min)
                sm.max = ifelse(merged_sm.max > sm.max, merged_sm.max, sm.max)
                sm.sum += merged_sm.sum
                sm.count += merged_sm.count
            else
                merged[city] = sm
            end
        end
    end
    merged
end

function process_chunk(data, chunk)
    stations = Dict{String,StationMeasurements}()
    for i in eachindex(chunk)
        if i == 1
            continue
        end
        line = String(data[chunk[i-1]:chunk[i]-1])
        station, temp_str = rsplit(line, ";")
        temp = parse(Float32, temp_str)
        if haskey(stations, station)
            sm = stations[station]
            sm.min = ifelse(temp < sm.min, temp, sm.min)
            sm.max = ifelse(temp > sm.max, temp, sm.max)
            sm.sum += temp
            sm.count += 1
        else
            stations[station] = StationMeasurements(temp, temp, temp, 1)
        end
    end
    stations
end

function main()
    open("../../../measurements.txt", "r") do f
        sz = Base.stat(f).size
        data = mmap(f, Vector{UInt8}, sz)
        idxs = findall(isequal(0x0a), data)
        idxs_chunks = collect(Iterators.partition(idxs, length(idxs) ÷ Threads.nthreads()))
        tasks = map(idxs_chunks) do chunk
            Threads.@spawn process_chunk(data, chunk)
        end
        stations_vec = fetch.(tasks)
        stations = merge(stations_vec)
        print_measurements(stations)
    end
end

main()