/*
 * Copyright (c) 2021, Hunter Salyer <thefalsehonesty@gmail.com>
 * Copyright (c) 2022, Gregory Bertilson <Zaggy1024@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#pragma once

#include <AK/IntegralMath.h>
#include <AK/Optional.h>
#include <AK/OwnPtr.h>
#include <LibCore/MappedFile.h>
#include <LibMedia/DecoderError.h>

#include "Document.h"

namespace Media::Matroska {

class SampleIterator;
class Streamer;

class Reader {
public:
    typedef Function<DecoderErrorOr<IterationDecision>(TrackEntry const&)> TrackEntryCallback;

    static DecoderErrorOr<Reader> from_file(StringView path);
    static DecoderErrorOr<Reader> from_mapped_file(NonnullOwnPtr<Core::MappedFile> mapped_file);

    static DecoderErrorOr<Reader> from_data(ReadonlyBytes data);

    EBMLHeader const& header() const { return m_header.value(); }

    DecoderErrorOr<SegmentInformation> segment_information();

    DecoderErrorOr<void> for_each_track(TrackEntryCallback);
    DecoderErrorOr<void> for_each_track_of_type(TrackEntry::TrackType, TrackEntryCallback);
    DecoderErrorOr<NonnullRefPtr<TrackEntry>> track_for_track_number(u64);
    DecoderErrorOr<size_t> track_count();

    DecoderErrorOr<SampleIterator> create_sample_iterator(u64 track_number);
    DecoderErrorOr<SampleIterator> seek_to_random_access_point(SampleIterator, Duration);
    DecoderErrorOr<Optional<Vector<CuePoint> const&>> cue_points_for_track(u64 track_number);
    DecoderErrorOr<bool> has_cues_for_track(u64 track_number);

private:
    Reader(ReadonlyBytes data)
        : m_data(data)
    {
    }

    DecoderErrorOr<void> parse_initial_data();

    DecoderErrorOr<Optional<size_t>> find_first_top_level_element_with_id([[maybe_unused]] StringView element_name, u32 element_id);

    DecoderErrorOr<void> ensure_tracks_are_parsed();
    DecoderErrorOr<void> parse_tracks(Streamer&);

    DecoderErrorOr<void> parse_cues(Streamer&);
    DecoderErrorOr<void> ensure_cues_are_parsed();
    DecoderErrorOr<void> seek_to_cue_for_timestamp(SampleIterator&, Duration const&);

    RefPtr<Core::SharedMappedFile> m_mapped_file;
    ReadonlyBytes m_data;

    Optional<EBMLHeader> m_header;

    size_t m_segment_contents_position { 0 };
    size_t m_segment_contents_size { 0 };

    HashMap<u32, size_t> m_seek_entries;
    size_t m_last_top_level_element_position { 0 };

    Optional<SegmentInformation> m_segment_information;

    OrderedHashMap<u64, NonnullRefPtr<TrackEntry>> m_tracks;

    // The vectors must be sorted by timestamp at all times.
    HashMap<u64, Vector<CuePoint>> m_cues;
    bool m_cues_have_been_parsed { false };
};

class SampleIterator {
public:
    DecoderErrorOr<Block> next_block();
    Cluster const& current_cluster() const { return *m_current_cluster; }
    Optional<Duration> const& last_timestamp() const { return m_last_timestamp; }
    TrackEntry const& track() const { return *m_track; }

private:
    friend class Reader;

    SampleIterator(RefPtr<Core::SharedMappedFile> file, ReadonlyBytes data, TrackEntry& track, u64 timestamp_scale, size_t position)
        : m_file(move(file))
        , m_data(data)
        , m_track(track)
        , m_segment_timestamp_scale(timestamp_scale)
        , m_position(position)
    {
    }

    DecoderErrorOr<void> seek_to_cue_point(CuePoint const& cue_point);

    RefPtr<Core::SharedMappedFile> m_file;
    ReadonlyBytes m_data;
    NonnullRefPtr<TrackEntry> m_track;
    u64 m_segment_timestamp_scale { 0 };

    // Must always point to an element ID or the end of the stream.
    size_t m_position { 0 };

    Optional<Duration> m_last_timestamp;

    Optional<Cluster> m_current_cluster;
};

class Streamer {
public:
    Streamer(ReadonlyBytes data)
        : m_data(data)
    {
    }

    u8 const* data() { return m_data.data() + m_position; }

    char const* data_as_chars() { return reinterpret_cast<char const*>(data()); }

    size_t octets_read() { return m_octets_read.last(); }

    void push_octets_read() { m_octets_read.append(0); }

    void pop_octets_read()
    {
        auto popped = m_octets_read.take_last();
        if (!m_octets_read.is_empty())
            m_octets_read.last() += popped;
    }

    ErrorOr<u8> read_octet();

    ErrorOr<i16> read_i16();

    ErrorOr<u64> read_variable_size_integer(bool mask_length = true);
    ErrorOr<i64> read_variable_size_signed_integer();

    ErrorOr<u64> read_u64();
    ErrorOr<double> read_float();

    ErrorOr<ByteString> read_string();

    ErrorOr<void> read_unknown_element();

    ErrorOr<ReadonlyBytes> read_raw_octets(size_t num_octets);

    size_t position() const { return m_position; }
    size_t remaining() const { return m_data.size() - position(); }

    bool at_end() const { return remaining() == 0; }
    bool has_octet() const { return remaining() >= 1; }

    ErrorOr<void> seek_to_position(size_t position);

private:
    ReadonlyBytes m_data;
    size_t m_position { 0 };
    Vector<size_t> m_octets_read { 0 };
};

}