#include "asset_manager.h"
#include "../Rendering/mesh.h"
#include "../Rendering/texture.h"
#include "../Rendering/shader.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#define ASSET_BUCKET_COUNT 256

struct AssetManager {
    Asset* buckets[ASSET_BUCKET_COUNT];
    uint32_t next_id;
};

static uint32_t hash_str(const char* s) {
    uint32_t h = 5381;
    while (*s) h = ((h << 5) + h) + (unsigned char)*s++;
    return h % ASSET_BUCKET_COUNT;
}

static int str_ends_with(const char* s, const char* suffix) {
    size_t ls = strlen(s), lf = strlen(suffix);
    return ls >= lf && strcmp(s + ls - lf, suffix) == 0;
}

static Asset* asset_alloc(AssetManager* am, AssetType type, const char* path, void* data, void (*destroy)(void*)) {
    Asset* a = (Asset*)calloc(1, sizeof(Asset));
    a->id = ++am->next_id;
    a->type = type;
    a->refcount = 1;
    a->owner = am;
    strncpy(a->path, path, sizeof(a->path) - 1);
    a->data = data;
    a->destroy = destroy;

    uint32_t h = hash_str(path);
    a->next = am->buckets[h];
    am->buckets[h] = a;
    return a;
}

static void asset_free(AssetManager* am, Asset* a) {
    if (!a) return;
    uint32_t h = hash_str(a->path);
    Asset** cur = &am->buckets[h];
    while (*cur) {
        if (*cur == a) {
            *cur = a->next;
            break;
        }
        cur = &(*cur)->next;
    }
    if (a->destroy) a->destroy(a->data);
    free(a);
}

AssetManager* asset_manager_create(void) {
    AssetManager* am = (AssetManager*)calloc(1, sizeof(AssetManager));
    return am;
}

void asset_manager_destroy(AssetManager* am) {
    if (!am) return;
    for (int i = 0; i < ASSET_BUCKET_COUNT; i++) {
        Asset* a = am->buckets[i];
        while (a) {
            Asset* next = a->next;
            if (a->destroy) a->destroy(a->data);
            free(a);
            a = next;
        }
    }
    free(am);
}

Asset* asset_manager_find(AssetManager* am, const char* path) {
    if (!am || !path) return NULL;
    uint32_t h = hash_str(path);
    for (Asset* a = am->buckets[h]; a; a = a->next) {
        if (strcmp(a->path, path) == 0) return a;
    }
    return NULL;
}

Asset* asset_manager_get(AssetManager* am, uint32_t id) {
    if (!am || id == 0) return NULL;
    for (int i = 0; i < ASSET_BUCKET_COUNT; i++) {
        for (Asset* a = am->buckets[i]; a; a = a->next) {
            if (a->id == id) return a;
        }
    }
    return NULL;
}

void* asset_data(Asset* a) { return a ? a->data : NULL; }

void asset_acquire(Asset* a) { if (a) a->refcount++; }

void asset_release(Asset* a) {
    if (!a) return;
    a->refcount--;
    if (a->refcount <= 0 && a->owner) {
        asset_free(a->owner, a);
    }
}

// ------------------------------------------------------------------
// Mesh
// ------------------------------------------------------------------
static void destroy_mesh_asset(void* data) { mesh_destroy((Mesh*)data); }

Asset* asset_manager_load_mesh(AssetManager* am, ID3D11Device* dev, const char* path) {
    if (!am || !dev || !path) return NULL;
    Asset* existing = asset_manager_find(am, path);
    if (existing) { asset_acquire(existing); return existing; }

    Mesh* mesh = NULL;
    if (strcmp(path, "cube") == 0) mesh = mesh_create_cube(dev);
    else if (strcmp(path, "sphere") == 0) mesh = mesh_create_sphere(dev, 1.0f, 16);
    else if (strcmp(path, "plane") == 0) mesh = mesh_create_plane(dev, 10.0f);
    else if (str_ends_with(path, ".obj")) mesh = mesh_load_obj(dev, path);

    if (!mesh) return NULL;
    return asset_alloc(am, ASSET_MESH, path, mesh, destroy_mesh_asset);
}

// ------------------------------------------------------------------
// Texture
// ------------------------------------------------------------------
static void destroy_texture_asset(void* data) { texture_destroy((Texture*)data); }

Asset* asset_manager_load_texture(AssetManager* am, ID3D11Device* dev, const char* path) {
    if (!am || !dev || !path) return NULL;
    Asset* existing = asset_manager_find(am, path);
    if (existing) { asset_acquire(existing); return existing; }

    Texture* tex = NULL;
    if (strcmp(path, "checkerboard") == 0) tex = texture_create_checkerboard(dev, 256, 8);
    else tex = texture_load_from_file(dev, path);

    if (!tex) return NULL;
    return asset_alloc(am, ASSET_TEXTURE, path, tex, destroy_texture_asset);
}

// ------------------------------------------------------------------
// Shader
// ------------------------------------------------------------------
static void destroy_shader_asset(void* data) { shader_destroy((Shader*)data); }

static char* read_file(const char* path, uint32_t* out_len) {
    FILE* f = fopen(path, "rb");
    if (!f) return NULL;
    fseek(f, 0, SEEK_END);
    long size = ftell(f);
    fseek(f, 0, SEEK_SET);
    if (size <= 0) { fclose(f); return NULL; }
    char* buf = (char*)malloc((size_t)size + 1);
    if (!buf) { fclose(f); return NULL; }
    fread(buf, 1, (size_t)size, f);
    buf[size] = '\0';
    fclose(f);
    if (out_len) *out_len = (uint32_t)size;
    return buf;
}

Asset* asset_manager_load_shader(AssetManager* am, ID3D11Device* dev, const char* vs_path, const char* ps_path) {
    if (!am || !dev) return NULL;

    char key[512];
    if (vs_path && ps_path) snprintf(key, sizeof(key), "%s|%s", vs_path, ps_path);
    else strcpy(key, "default");

    Asset* existing = asset_manager_find(am, key);
    if (existing) { asset_acquire(existing); return existing; }

    Shader* s = shader_create();
    if (vs_path && ps_path) {
        uint32_t vs_len = 0, ps_len = 0;
        char* vs = read_file(vs_path, &vs_len);
        char* ps = read_file(ps_path, &ps_len);
        if (!vs || !ps || !shader_compile(s, dev, vs, ps, vs_len, ps_len)) {
            shader_destroy(s); free(vs); free(ps); return NULL;
        }
        free(vs); free(ps);
    } else {
        if (!shader_compile(s, dev, NULL, NULL, 0, 0)) {
            shader_destroy(s); return NULL;
        }
    }
    return asset_alloc(am, ASSET_SHADER, key, s, destroy_shader_asset);
}