#include "faststart.h"

#include <android/log.h>

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <vector>

#define TAG "MoovFastStart"
#define LOGI(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__)

static uint32_t ReadBE32(FILE* f) {
    uint8_t buf[4];
    if (fread(buf, 1, 4, f) != 4) return 0;
    return (uint32_t(buf[0]) << 24) | (uint32_t(buf[1]) << 16) |
           (uint32_t(buf[2]) << 8) | uint32_t(buf[3]);
}

static void WriteBE32(uint8_t* p, uint32_t v) {
    p[0] = (v >> 24) & 0xFF;
    p[1] = (v >> 16) & 0xFF;
    p[2] = (v >> 8) & 0xFF;
    p[3] = v & 0xFF;
}

static void WriteBE64(uint8_t* p, uint64_t v) {
    WriteBE32(p, (uint32_t)(v >> 32));
    WriteBE32(p + 4, (uint32_t)(v & 0xFFFFFFFF));
}

static uint64_t ReadBE64(const uint8_t* p) {
    return ((uint64_t)((uint32_t(p[0]) << 24) | (uint32_t(p[1]) << 16) |
                       (uint32_t(p[2]) << 8) | uint32_t(p[3])) << 32) |
           ((uint32_t(p[4]) << 24) | (uint32_t(p[5]) << 16) |
            (uint32_t(p[6]) << 8) | uint32_t(p[7]));
}

// Adjust chunk offsets in moov by a delta (moov size)
static void AdjustChunkOffsets(uint8_t* data, uint32_t size, int64_t delta) {
    if (size < 8) return;

    uint32_t atomSize = (uint32_t(data[0]) << 24) | (uint32_t(data[1]) << 16) |
                        (uint32_t(data[2]) << 8) | uint32_t(data[3]);
    char fourcc[5] = {(char)data[4], (char)data[5], (char)data[6], (char)data[7], 0};

    if (atomSize == 0 || atomSize > size) atomSize = size;

    if (strcmp(fourcc, "stco") == 0 && atomSize >= 16) {
        // 32-bit chunk offset box: 4 size + 4 fourcc + 1 version + 3 flags + 4 count + N*4 offsets
        uint32_t count = (uint32_t(data[12]) << 24) | (uint32_t(data[13]) << 16) |
                         (uint32_t(data[14]) << 8) | uint32_t(data[15]);
        for (uint32_t i = 0; i < count && (16 + (i + 1) * 4) <= atomSize; i++) {
            uint8_t* p = data + 16 + i * 4;
            uint32_t offset = (uint32_t(p[0]) << 24) | (uint32_t(p[1]) << 16) |
                              (uint32_t(p[2]) << 8) | uint32_t(p[3]);
            WriteBE32(p, (uint32_t)(offset + delta));
        }
    } else if (strcmp(fourcc, "co64") == 0 && atomSize >= 16) {
        // 64-bit chunk offset box
        uint32_t count = (uint32_t(data[12]) << 24) | (uint32_t(data[13]) << 16) |
                         (uint32_t(data[14]) << 8) | uint32_t(data[15]);
        for (uint32_t i = 0; i < count && (16 + (i + 1) * 8) <= atomSize; i++) {
            uint8_t* p = data + 16 + i * 8;
            uint64_t offset = ReadBE64(p);
            WriteBE64(p, offset + delta);
        }
    } else {
        // Container atom — recurse into children
        // Skip header (8 bytes for regular, 16 for version+flags atoms)
        bool isContainer = (strcmp(fourcc, "moov") == 0 || strcmp(fourcc, "trak") == 0 ||
                            strcmp(fourcc, "mdia") == 0 || strcmp(fourcc, "minf") == 0 ||
                            strcmp(fourcc, "stbl") == 0 || strcmp(fourcc, "edts") == 0 ||
                            strcmp(fourcc, "dinf") == 0 || strcmp(fourcc, "udta") == 0);
        if (isContainer) {
            uint32_t offset = 8;
            while (offset + 8 <= atomSize) {
                uint32_t childSize = (uint32_t(data[offset]) << 24) | (uint32_t(data[offset+1]) << 16) |
                                     (uint32_t(data[offset+2]) << 8) | uint32_t(data[offset+3]);
                if (childSize < 8 || offset + childSize > atomSize) break;
                AdjustChunkOffsets(data + offset, childSize, delta);
                offset += childSize;
            }
        }
    }
}

static bool CopyBytes(FILE* src, FILE* dst, int64_t count) {
    uint8_t buf[8192];
    while (count > 0) {
        size_t toRead = (count > (int64_t)sizeof(buf)) ? sizeof(buf) : (size_t)count;
        size_t read = fread(buf, 1, toRead, src);
        if (read != toRead) return false;
        if (fwrite(buf, 1, read, dst) != read) return false;
        count -= read;
    }
    return true;
}

bool MoovFastStart(const std::string& inputPath, const std::string& outputPath) {
    FILE* input = fopen(inputPath.c_str(), "rb");
    if (!input) {
        LOGE("Cannot open input: %s", inputPath.c_str());
        return false;
    }

    // Scan for atoms
    struct Atom { char fourcc[5]; int64_t offset; int64_t size; };
    std::vector<Atom> atoms;

    fseek(input, 0, SEEK_END);
    int64_t fileSize = ftell(input);
    fseek(input, 0, SEEK_SET);

    int64_t pos = 0;
    while (pos < fileSize) {
        fseek(input, pos, SEEK_SET);
        uint32_t size32 = ReadBE32(input);
        uint32_t fourcc_raw = ReadBE32(input);

        Atom atom;
        atom.fourcc[0] = (fourcc_raw >> 24) & 0xFF;
        atom.fourcc[1] = (fourcc_raw >> 16) & 0xFF;
        atom.fourcc[2] = (fourcc_raw >> 8) & 0xFF;
        atom.fourcc[3] = fourcc_raw & 0xFF;
        atom.fourcc[4] = 0;
        atom.offset = pos;

        if (size32 == 1) {
            // 64-bit extended size
            uint8_t ext[8];
            if (fread(ext, 1, 8, input) != 8) break;
            atom.size = (int64_t)ReadBE64(ext);
        } else if (size32 == 0) {
            atom.size = fileSize - pos;
        } else {
            atom.size = size32;
        }

        if (atom.size < 8) break;
        atoms.push_back(atom);
        pos += atom.size;
    }

    // Find moov and mdat
    int moovIdx = -1, mdatIdx = -1;
    for (int i = 0; i < (int)atoms.size(); i++) {
        if (strcmp(atoms[i].fourcc, "moov") == 0) moovIdx = i;
        if (strcmp(atoms[i].fourcc, "mdat") == 0) mdatIdx = i;
    }

    if (moovIdx < 0) {
        LOGE("No moov atom found");
        fclose(input);
        return false;
    }

    if (moovIdx < mdatIdx || mdatIdx < 0) {
        // moov already before mdat — just copy as-is
        LOGI("moov already at front, copying file");
        fclose(input);
        if (inputPath != outputPath) {
            FILE* in2 = fopen(inputPath.c_str(), "rb");
            FILE* out = fopen(outputPath.c_str(), "wb");
            if (!in2 || !out) {
                if (in2) fclose(in2);
                if (out) fclose(out);
                return false;
            }
            bool ok = CopyBytes(in2, out, fileSize);
            fclose(in2);
            fclose(out);
            return ok;
        }
        return true;
    }

    // Read moov into memory
    int64_t moovSize = atoms[moovIdx].size;
    std::vector<uint8_t> moovData(moovSize);
    fseek(input, atoms[moovIdx].offset, SEEK_SET);
    if (fread(moovData.data(), 1, moovSize, input) != (size_t)moovSize) {
        LOGE("Failed to read moov atom");
        fclose(input);
        return false;
    }

    // Adjust chunk offsets by moov size (since moov will be inserted before mdat)
    AdjustChunkOffsets(moovData.data(), moovData.size(), moovSize);

    // Write output: atoms before mdat + moov + mdat + atoms after moov
    FILE* output = fopen(outputPath.c_str(), "wb");
    if (!output) {
        LOGE("Cannot open output: %s", outputPath.c_str());
        fclose(input);
        return false;
    }

    // Write all atoms before mdat
    for (int i = 0; i < mdatIdx; i++) {
        fseek(input, atoms[i].offset, SEEK_SET);
        if (!CopyBytes(input, output, atoms[i].size)) {
            LOGE("Failed to copy pre-mdat atom");
            fclose(input);
            fclose(output);
            return false;
        }
    }

    // Write moov (with adjusted offsets)
    if (fwrite(moovData.data(), 1, moovSize, output) != (size_t)moovSize) {
        LOGE("Failed to write moov");
        fclose(input);
        fclose(output);
        return false;
    }

    // Write mdat and everything after (except moov which we already wrote)
    for (int i = mdatIdx; i < (int)atoms.size(); i++) {
        if (i == moovIdx) continue; // skip moov — already written
        fseek(input, atoms[i].offset, SEEK_SET);
        if (!CopyBytes(input, output, atoms[i].size)) {
            LOGE("Failed to copy post-moov atom");
            fclose(input);
            fclose(output);
            return false;
        }
    }

    fclose(input);
    fclose(output);
    LOGI("Faststart complete: moov moved to front");
    return true;
}
