#include "jail_shader.h"

#ifdef __APPLE__
#include <SDL3/SDL.h>
#include <SDL3/SDL_metal.h>
#include <Metal/Metal.h>
#include <QuartzCore/CAMetalLayer.h>
#include <CoreFoundation/CoreFoundation.h>
#include <stdexcept>
#include <vector>
#include "../asset.h"

namespace shader {
namespace metal {

// Metal objects
id<MTLDevice> device = nullptr;
id<MTLCommandQueue> commandQueue = nullptr;
id<MTLRenderPipelineState> pipelineState = nullptr;
id<MTLTexture> metalTexture = nullptr;
id<MTLBuffer> vertexBuffer = nullptr;
CAMetalLayer* metalLayer = nullptr;
SDL_MetalView metalView = nullptr;

// Vertex data for fullscreen quad
struct Vertex {
    float position[4];  // x, y, z, w
    float texcoord[2];  // u, v
};

const Vertex quadVertices[] = {
    // Position (x, y, z, w)     // TexCoord (u, v)
    {{-1.0f, -1.0f, 0.0f, 1.0f}, {0.0f, 1.0f}}, // Bottom-left
    {{ 1.0f, -1.0f, 0.0f, 1.0f}, {1.0f, 1.0f}}, // Bottom-right
    {{-1.0f,  1.0f, 0.0f, 1.0f}, {0.0f, 0.0f}}, // Top-left
    {{ 1.0f,  1.0f, 0.0f, 1.0f}, {1.0f, 0.0f}}, // Top-right
};

std::string loadMetalShader(const std::string& filename) {
    // Try to load the .metal file from the same location as GLSL files
    auto data = Asset::get()->loadData(filename);
    if (!data.empty()) {
        return std::string(data.begin(), data.end());
    }
    return "";
}

bool initMetal(SDL_Window* window, SDL_Texture* backBuffer, const std::string& shaderFilename) {
    // Get Metal view from SDL
    metalView = SDL_Metal_CreateView(window);
    if (!metalView) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to create Metal view: %s", SDL_GetError());
        return false;
    }
    
    metalLayer = (__bridge CAMetalLayer*)SDL_Metal_GetLayer(metalView);
    if (!metalLayer) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to get Metal layer");
        return false;
    }
    
    // Get Metal device
    device = MTLCreateSystemDefaultDevice();
    if (!device) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to create Metal device");
        return false;
    }
    
    // Configure the layer
    metalLayer.device = device;
    metalLayer.pixelFormat = MTLPixelFormatBGRA8Unorm;
    metalLayer.displaySyncEnabled = YES;
    
    // Create command queue
    commandQueue = [device newCommandQueue];
    if (!commandQueue) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to create Metal command queue");
        return false;
    }
    
    // Load and compile shaders
    std::string metalShaderSource = loadMetalShader(shaderFilename);
    if (metalShaderSource.empty()) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to load Metal shader: %s", shaderFilename.c_str());
        return false;
    }
    
    SDL_LogInfo(SDL_LOG_CATEGORY_APPLICATION, "Loaded Metal shader %s (length: %zu)", 
                shaderFilename.c_str(), metalShaderSource.length());
    
    NSString* shaderNSString = [NSString stringWithUTF8String:metalShaderSource.c_str()];
    NSError* error = nil;
    id<MTLLibrary> library = [device newLibraryWithSource:shaderNSString options:nil error:&error];
    if (!library || error) {
        if (error) {
            SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to compile Metal shader: %s", 
                        [[error localizedDescription] UTF8String]);
        }
        return false;
    }
    
    id<MTLFunction> vertexFunction = [library newFunctionWithName:@"vertex_main"];
    id<MTLFunction> fragmentFunction = [library newFunctionWithName:@"fragment_main"];
    
    if (!vertexFunction || !fragmentFunction) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to load Metal shader functions");
        return false;
    }
    
    // Create render pipeline
    MTLRenderPipelineDescriptor* pipelineDescriptor = [[MTLRenderPipelineDescriptor alloc] init];
    pipelineDescriptor.vertexFunction = vertexFunction;
    pipelineDescriptor.fragmentFunction = fragmentFunction;
    pipelineDescriptor.colorAttachments[0].pixelFormat = metalLayer.pixelFormat;
    
    // Set up vertex descriptor
    MTLVertexDescriptor* vertexDescriptor = [[MTLVertexDescriptor alloc] init];
    vertexDescriptor.attributes[0].format = MTLVertexFormatFloat4;
    vertexDescriptor.attributes[0].offset = 0;
    vertexDescriptor.attributes[0].bufferIndex = 0;
    vertexDescriptor.attributes[1].format = MTLVertexFormatFloat2;
    vertexDescriptor.attributes[1].offset = 4 * sizeof(float);
    vertexDescriptor.attributes[1].bufferIndex = 0;
    vertexDescriptor.layouts[0].stride = sizeof(Vertex);
    vertexDescriptor.layouts[0].stepFunction = MTLVertexStepFunctionPerVertex;
    
    pipelineDescriptor.vertexDescriptor = vertexDescriptor;
    
    pipelineState = [device newRenderPipelineStateWithDescriptor:pipelineDescriptor error:&error];
    if (!pipelineState || error) {
        if (error) {
            SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to create Metal render pipeline: %s", 
                        [[error localizedDescription] UTF8String]);
        }
        return false;
    }
    
    // Create vertex buffer
    vertexBuffer = [device newBufferWithBytes:quadVertices 
                                       length:sizeof(quadVertices) 
                                      options:MTLResourceOptionCPUCacheModeDefault];
    
    SDL_LogInfo(SDL_LOG_CATEGORY_APPLICATION, "** Metal shader system initialized successfully");
    return true;
}

void updateMetalTexture(SDL_Texture* backBuffer) {
    if (!device || !backBuffer) return;
    
    // Solo intentar obtener la textura una vez al inicio, después reutilizar
    static bool textureInitialized = false;
    if (textureInitialized && metalTexture) {
        return; // Ya tenemos la textura configurada
    }
    
    // Intentar obtener la textura Metal directamente desde SDL
    SDL_PropertiesID props = SDL_GetTextureProperties(backBuffer);
    if (props == 0) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to get texture properties: %s", SDL_GetError());
        return;
    }
    
    // Obtener la textura Metal nativa de SDL
    id<MTLTexture> sdlMetalTexture = (__bridge id<MTLTexture>)SDL_GetPointerProperty(props, "SDL.texture.metal.texture", nullptr);
    
    if (sdlMetalTexture) {
        // Usar la textura Metal directamente de SDL
        metalTexture = sdlMetalTexture;
        textureInitialized = true;
        SDL_LogInfo(SDL_LOG_CATEGORY_APPLICATION, "Using SDL Metal texture directly (size: %lux%lu)", 
                   [metalTexture width], [metalTexture height]);
        return;
    }
    
    // Intentar con nombres alternativos
    const char* propertyNames[] = {
        "texture.metal.texture",
        "SDL.renderer.metal.texture",
        "metal.texture"
    };
    
    for (const char* propName : propertyNames) {
        sdlMetalTexture = (__bridge id<MTLTexture>)SDL_GetPointerProperty(props, propName, nullptr);
        if (sdlMetalTexture) {
            metalTexture = sdlMetalTexture;
            textureInitialized = true;
            SDL_LogInfo(SDL_LOG_CATEGORY_APPLICATION, "Using SDL Metal texture via property '%s'", propName);
            return;
        }
    }
    
    // Si no podemos obtener la textura Metal, crear una propia que coincida con la de SDL
    if (!metalTexture) {
        float width_f, height_f;
        if (SDL_GetTextureSize(backBuffer, &width_f, &height_f)) {
            MTLTextureDescriptor* desc = [MTLTextureDescriptor texture2DDescriptorWithPixelFormat:MTLPixelFormatBGRA8Unorm
                                                                                                width:(int)width_f
                                                                                               height:(int)height_f
                                                                                            mipmapped:NO];
            desc.usage = MTLTextureUsageShaderRead | MTLTextureUsageRenderTarget;
            metalTexture = [device newTextureWithDescriptor:desc];
            textureInitialized = true;
            SDL_LogWarn(SDL_LOG_CATEGORY_APPLICATION, "Created fallback Metal texture (%dx%d)", 
                       (int)width_f, (int)height_f);
        }
    }
}

void renderMetal() {
    if (!device || !commandQueue || !pipelineState || !metalLayer) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Metal render failed: missing components");
        return;
    }
    
    // Get the current drawable
    id<CAMetalDrawable> drawable = [metalLayer nextDrawable];
    if (!drawable) {
        SDL_LogError(SDL_LOG_CATEGORY_APPLICATION, "Failed to get Metal drawable");
        return;
    }
    
    // Create render pass descriptor
    MTLRenderPassDescriptor* renderPass = [MTLRenderPassDescriptor renderPassDescriptor];
    renderPass.colorAttachments[0].texture = drawable.texture;
    renderPass.colorAttachments[0].loadAction = MTLLoadActionClear;
    renderPass.colorAttachments[0].clearColor = MTLClearColorMake(0.0, 0.0, 0.0, 1.0);
    renderPass.colorAttachments[0].storeAction = MTLStoreActionStore;
    
    // Create command buffer
    id<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];
    
    // Create render command encoder
    id<MTLRenderCommandEncoder> encoder = [commandBuffer renderCommandEncoderWithDescriptor:renderPass];
    
    // Set pipeline state
    [encoder setRenderPipelineState:pipelineState];
    
    // Set vertex buffer
    [encoder setVertexBuffer:vertexBuffer offset:0 atIndex:0];
    
    // Set texture and sampler (use metalTexture if available, otherwise render solid color)
    if (metalTexture) {
        [encoder setFragmentTexture:metalTexture atIndex:0];
        
        // Create sampler state for nearest neighbor filtering (pixelated look)
        MTLSamplerDescriptor* samplerDescriptor = [[MTLSamplerDescriptor alloc] init];
        samplerDescriptor.minFilter = MTLSamplerMinMagFilterNearest;
        samplerDescriptor.magFilter = MTLSamplerMinMagFilterNearest;
        samplerDescriptor.sAddressMode = MTLSamplerAddressModeClampToEdge;
        samplerDescriptor.tAddressMode = MTLSamplerAddressModeClampToEdge;
        
        id<MTLSamplerState> sampler = [device newSamplerStateWithDescriptor:samplerDescriptor];
        [encoder setFragmentSamplerState:sampler atIndex:0];
    } else {
        SDL_LogInfo(SDL_LOG_CATEGORY_APPLICATION, "No Metal texture - rendering test pattern");
        // No texture available - the shader should handle this case
    }
    
    // Draw the quad
    [encoder drawPrimitives:MTLPrimitiveTypeTriangleStrip vertexStart:0 vertexCount:4];
    
    // End encoding
    [encoder endEncoding];
    
    // Present drawable
    [commandBuffer presentDrawable:drawable];
    
    // Commit command buffer
    [commandBuffer commit];
}

void cleanupMetal() {
    // Release Metal objects (ARC handles most of this automatically)
    pipelineState = nullptr;
    metalTexture = nullptr;
    vertexBuffer = nullptr;
    commandQueue = nullptr;
    device = nullptr;
    
    if (metalView) {
        SDL_Metal_DestroyView(metalView);
        metalView = nullptr;
    }
    
    metalLayer = nullptr;
    
    SDL_LogInfo(SDL_LOG_CATEGORY_APPLICATION, "Metal shader system cleaned up");
}

} // namespace metal
} // namespace shader

#endif // __APPLE__