#include <kos.h>
#include <stdio.h>
#include <dirent.h>
#include <dc/video.h>
#include <dc/biosfont.h>

#define PL_MPEG_IMPLEMENTATION
#include "pl_mpeg.h"

static pvr_ptr_t texture_ptrs[2];
static int current_texture = 0;
static int video_width = 0;
static int video_height = 0;
static const int tex_size = 512; 

void draw_error(const char *msg) {
    uint16_t *vram = vram_s;
    int x = 20, y = 20;
    for(int i = 0; i < 640 * 480; i++) vram[i] = 0x000F;
    bfont_draw_str(vram + y * 640 + x, 640, 0, msg);
    while(1) {
        MAPLE_FOREACH_BEGIN(MAPLE_FUNC_CONTROLLER, cont_state_t, st)
            if(st->buttons & CONT_START) return;
        MAPLE_FOREACH_END()
    }
}

void draw_frame() {
    pvr_poly_cxt_t cxt;
    pvr_poly_hdr_t hdr;
    pvr_vertex_t vert;

    pvr_wait_ready();
    pvr_scene_begin();
    pvr_list_begin(PVR_LIST_OP_POLY);
    pvr_poly_cxt_txr(&cxt, PVR_LIST_OP_POLY, PVR_TXRFMT_YUV422 | PVR_TXRFMT_NONTWIDDLED, 
                     tex_size, tex_size, texture_ptrs[current_texture], PVR_FILTER_BILINEAR);
    pvr_poly_compile(&hdr, &cxt);
    pvr_prim(&hdr, sizeof(hdr));

    float u_max = (float)video_width / (float)tex_size;
    float v_max = (float)video_height / (float)tex_size;

    vert.flags = PVR_CMD_VERTEX;
    vert.z = 1.0f;
    vert.argb = PVR_PACK_COLOR(1.0f, 1.0f, 1.0f, 1.0f);

    vert.x = 0; vert.y = 0; vert.u = 0; vert.v = 0;
    pvr_prim(&vert, sizeof(vert));

    vert.x = 640; vert.y = 0; vert.u = u_max; vert.v = 0;
    pvr_prim(&vert, sizeof(vert));

    vert.x = 0; vert.y = 480; vert.u = 0; vert.v = v_max;
    pvr_prim(&vert, sizeof(vert));

    vert.flags = PVR_CMD_VERTEX_EOL;
    vert.x = 640; vert.y = 480; vert.u = u_max; vert.v = v_max;
    pvr_prim(&vert, sizeof(vert));

    pvr_list_finish();
    pvr_scene_finish();
}


void list_files_on_screen() {
    DIR *dir;
    struct dirent *ent;
    int x = 20, y = 50;
    uint16_t *vram = vram_s;

    dir = opendir("/cd");
    if (dir == NULL) {
        bfont_draw_str(vram + (y * 640) + x, 640, 0, "Error: Could not open /cd");
        return;
    }

    bfont_draw_str(vram + (y * 640) + x, 640, 0, "Files found on /cd:");
    y += 20;

    while ((ent = readdir(dir)) != NULL) {
        char msg[128];
        sprintf(msg, "  -> %s", ent->d_name);
        bfont_draw_str(vram + (y * 640) + x, 640, 0, msg);
        y += 20;
        
        if (y > 450) break;
    }
    closedir(dir);
}

void convert_frame_yuv(plm_frame_t *frame, uint32_t vram_addr) {
    int mb_w = video_width / 16;
    int mb_h = video_height / 16;
    int tex_mb_w = tex_size / 16;
    uint8_t *y_plane = frame->y.data;
    uint8_t *u_plane = frame->cb.data;
    uint8_t *v_plane = frame->cr.data;

    // Set the YUV converter destination
    PVR_SET(PVR_YUV_ADDR, vram_addr & 0xffffff);
    // Config the YUV converter for YUV420 and 512x512 texture area
    PVR_SET(PVR_YUV_CFG, (0x00 << 24) | (((tex_size / 16) - 1) << 8) | ((tex_size / 16) - 1));
    // Required read
    PVR_GET(PVR_YUV_CFG);

    // Store Queue setup
    uint32_t *sq = sq_lock((void *)PVR_TA_YUV_CONV);
    uint32_t *u_block = (uint32_t *)SQ_MASK_DEST_ADDR(PVR_TA_YUV_CONV);
    uint32_t *v_block = (uint32_t *)SQ_MASK_DEST_ADDR(PVR_TA_YUV_CONV + 64);
    uint32_t *y_block = (uint32_t *)SQ_MASK_DEST_ADDR(PVR_TA_YUV_CONV + 128);

    for (int mby = 0; mby < mb_h; mby++) {
        for (int mbx = 0; mbx < mb_w; mbx++) {
            // PVR expects U, V, then 4 Y blocks
            
            // Send U block (8x8)
            uint8_t *src_u = u_plane + (mby * 8) * (video_width / 2) + (mbx * 8);
            for (int i = 0; i < 8; i++) {
                uint64_t *s = (uint64_t *)(src_u + i * (video_width / 2));
                *(uint64_t *)&((uint8_t *)u_block)[i * 8] = *s;
                if (!((i + 1) & 3)) sq_flush(&u_block[i * 2]);
            }

            // Send V block (8x8)
            uint8_t *src_v = v_plane + (mby * 8) * (video_width / 2) + (mbx * 8);
            for (int i = 0; i < 8; i++) {
                uint64_t *s = (uint64_t *)(src_v + i * (video_width / 2));
                *(uint64_t *)&((uint8_t *)v_block)[i * 8] = *s;
                if (!((i + 1) & 3)) sq_flush(&v_block[i * 2]);
            }

            // Send 4 Y blocks (8x8 each)
            for (int i = 0; i < 4; i++) {
                for (int j = 0; j < 8; j++) {
                    int y_off = (mby * 16 + j + (i / 2 * 8));
                    int x_off = mbx * 16 + (i % 2 * 8);
                    uint64_t *s = (uint64_t *)(y_plane + y_off * video_width + x_off);
                    *(uint64_t *)&((uint8_t *)y_block)[i * 64 + j * 8] = *s;
                    if (!((j + 1) & 3)) sq_flush(&y_block[(i * 64 + j * 8) / 4]);
                }
            }
        }
        
        // Pad the rest of the row with dummy macroblocks (if needed)
        int dummies = tex_mb_w - mb_w;
        if (dummies > 0) {
            for (int d = 0; d < dummies; d++) {
                // Send 12 SQ flushes of zeros for one dummy macroblock (384 bytes)
                for (int i = 0; i < 12; i++) {
                    sq[0] = sq[1] = sq[2] = sq[3] = sq[4] = sq[5] = sq[6] = sq[7] = 0;
                    sq_flush(sq);
                }
            }
        }
    }
    sq_unlock();
}

void play_video(const char *filename) {
    plm_t *plm = plm_create_with_filename(filename);
    if (!plm) {
        char err[128];
        sprintf(err, "FATAL: Could not open %s", filename);
        draw_error(err);
        return;
    }

    video_width = plm_get_width(plm);
    video_height = plm_get_height(plm);
    
    // Safety check for resolution
    if (video_width % 16 != 0 || video_height % 16 != 0) {
        draw_error("Resolution must be multiple of 16");
        plm_destroy(plm);
        return;
    }

    uint64_t start_time = timer_ms_gettime64();

    while (!plm_has_ended(plm)) {
        plm_frame_t *frame = plm_decode_video(plm);

        if (frame) {
            // Wait for system time to catch up to the current video time
            double video_time_ms = frame->time * 1000.0;
            while ((timer_ms_gettime64() - start_time) < video_time_ms) {
                thd_pass();
            }

            // Flip buffer
            current_texture = !current_texture;
            
            // Convert and transfer YUV macroblocks to VRAM
            convert_frame_yuv(frame, (uint32_t)texture_ptrs[current_texture]);
            
            draw_frame();
        }

        MAPLE_FOREACH_BEGIN(MAPLE_FUNC_CONTROLLER, cont_state_t, st)
            if(st->buttons & CONT_START) return;
        MAPLE_FOREACH_END()
    }
    plm_destroy(plm);
}

int main(int argc, char **argv) {
    pvr_init_defaults();
    texture_ptrs[0] = pvr_mem_malloc(tex_size * tex_size * 2);
    texture_ptrs[1] = pvr_mem_malloc(tex_size * tex_size * 2);
    
    if (!texture_ptrs[0] || !texture_ptrs[1]) return -1;

    play_video("/cd/VIDEO.MPG");

    pvr_mem_free(texture_ptrs[0]);
    pvr_mem_free(texture_ptrs[1]);
    return 0;
}