import { Canvas, loadImage, Image } from "canvas";
import { ReplayData } from "../types/replay";
import * as path from "path";

type ViewMode = 'top' | 'side' | 'front';

interface HeatmapConfig {
    width: number;
    height: number;
    worldSize: number;
    sensitivity: number;
    blurAmount: number;
    viewMode: ViewMode;
    backgroundImage?: string;
    preloadedBackground?: Image;
}

const defaultConfig: HeatmapConfig = {
    width: 100,
    height: 100,
    worldSize: 500,
    sensitivity: 0.7,
    blurAmount: 3,
    viewMode: 'top'
};

export const drawHeatmap = (canvas: Canvas, replayData: ReplayData, config: Partial<HeatmapConfig> = {}) => {
    const finalConfig = { ...defaultConfig, ...config };
    
    // Initialize heatmap array
    const heatmap: number[][] = Array(finalConfig.height).fill(null).map(() => Array(finalConfig.width).fill(0));
    
    let good = 0, bad = 0;
    
    // Process all player positions
    for (const player of replayData.players) {
        for (const snapshot of player.snapshots) {
            const pos = snapshot.position;
            
            // Convert world coordinates to heatmap coordinates based on view mode
            let x: number, y: number;
            
            switch (finalConfig.viewMode) {
                case 'top':
                    // Top view: X,Z -> heatmap X,Y
                    x = Math.floor((pos.x / finalConfig.worldSize + 0.5) * finalConfig.width);
                    y = Math.floor((pos.z / finalConfig.worldSize + 0.5) * finalConfig.height);
                    break;
                case 'side':
                    // Side view: Z,Y -> heatmap X,Y
                    x = Math.floor((pos.z / finalConfig.worldSize + 0.5) * finalConfig.width);
                    y = Math.floor((pos.y / finalConfig.worldSize + 0.5) * finalConfig.height);
                    break;
                case 'front':
                    // Front view: X,Y -> heatmap X,Y
                    x = Math.floor((pos.x / finalConfig.worldSize + 0.5) * finalConfig.width);
                    y = Math.floor((pos.y / finalConfig.worldSize + 0.5) * finalConfig.height);
                    break;
                default:
                    throw new Error(`Unknown view mode: ${finalConfig.viewMode}`);
            }
            
            if (x >= 0 && x < finalConfig.width && y >= 0 && y < finalConfig.height) {
                // Apply blur effect for smoother heatmap areas
                const blurRadius = finalConfig.blurAmount;
                for (let dx = -blurRadius; dx <= blurRadius; dx++) {
                    for (let dy = -blurRadius; dy <= blurRadius; dy++) {
                        const nx = x + dx;
                        const ny = y + dy;
                        
                        if (nx >= 0 && nx < finalConfig.width && ny >= 0 && ny < finalConfig.height) {
                            const distance = Math.sqrt(dx * dx + dy * dy);
                            // Gaussian-like falloff
                            const weight = Math.exp(-distance * distance / (2 * blurRadius * blurRadius));
                            heatmap[ny][nx] += Math.round(weight * 5);
                        }
                    }
                }
                good++;
            } else {
                bad++;
            }
        }
    }
    
    console.log(`Good: ${good}, Bad: ${bad}`);
    
    // Find maximum count for normalization
    let maxCount = 0;
    for (let y = 0; y < finalConfig.height; y++) {
        for (let x = 0; x < finalConfig.width; x++) {
            maxCount = Math.max(maxCount, heatmap[y][x]);
        }
    }
    
    // Create image data
    const ctx = canvas.getContext("2d");
    const imageData = ctx.createImageData(finalConfig.width, finalConfig.height);
    
    // Generate heatmap colors
    for (let y = 0; y < finalConfig.height; y++) {
        for (let x = 0; x < finalConfig.width; x++) {
            let value = heatmap[y][x] / maxCount;
            
            // Apply power curve to increase sensitivity to lower values
            value = Math.pow(value, finalConfig.sensitivity);
            
            const color = getHeatmapColor(value);
            const index = (y * finalConfig.width + x) * 4;
            
            imageData.data[index] = color.r;     // Red
            imageData.data[index + 1] = color.g; // Green
            imageData.data[index + 2] = color.b; // Blue
            imageData.data[index + 3] = 255;     // Alpha
        }
    }
    
    // Scale the image to canvas size
    const tempCanvas = canvas.constructor as any;
    const tempCtx = new tempCanvas(finalConfig.width, finalConfig.height).getContext("2d");
    tempCtx.putImageData(imageData, 0, 0);
    
    // Clear main canvas and draw scaled heatmap
    ctx.fillStyle = "white";
    ctx.fillRect(0, 0, canvas.width, canvas.height);
    ctx.drawImage(tempCtx.canvas, 0, 0, canvas.width, canvas.height);
};

function getHeatmapColor(value: number): { r: number; g: number; b: number } {
    // Blue -> Green -> Yellow -> Red gradient
    if (value < 0.33) {
        // Blue to Green (0.0 to 0.33)
        const t = value / 0.33;
        return {
            r: Math.round(lerp(0, 0, t) * 255),
            g: Math.round(lerp(0, 1, t) * 255),
            b: Math.round(lerp(1, 0, t) * 255)
        };
    } else if (value < 0.66) {
        // Green to Yellow (0.33 to 0.66)
        const t = (value - 0.33) / 0.33;
        return {
            r: Math.round(lerp(0, 1, t) * 255),
            g: Math.round(lerp(1, 1, t) * 255),
            b: Math.round(lerp(0, 0, t) * 255)
        };
    } else {
        // Yellow to Red (0.66 to 1.0)
        const t = (value - 0.66) / 0.34;
        return {
            r: Math.round(lerp(1, 1, t) * 255),
            g: Math.round(lerp(1, 0, t) * 255),
            b: Math.round(lerp(0, 0, t) * 255)
        };
    }
}

function lerp(a: number, b: number, t: number): number {
    return a + (b - a) * t;
}

export const drawTimeBasedHeatmap = async (canvas: Canvas, replayData: ReplayData, maxTime: number, config: Partial<HeatmapConfig> = {}) => {
    const finalConfig = { ...defaultConfig, ...config };
    
    // Initialize heatmap array
    const heatmap: number[][] = Array(finalConfig.height).fill(null).map(() => Array(finalConfig.width).fill(0));
    
    let good = 0, bad = 0;
    
    // Process all player positions up to maxTime
    for (const player of replayData.players) {
        for (const snapshot of player.snapshots) {
            // Only include positions up to the specified time
            if (snapshot.time > maxTime) {
                continue;
            }
            
            const pos = snapshot.position;
            
            // Convert world coordinates to heatmap coordinates based on view mode
            let x: number, y: number;
            
            switch (finalConfig.viewMode) {
                case 'top':
                    // Top view: X,Z -> heatmap X,Y
                    x = Math.floor((pos.x / finalConfig.worldSize + 0.5) * finalConfig.width);
                    y = Math.floor((pos.z / finalConfig.worldSize + 0.5) * finalConfig.height);
                    break;
                case 'side':
                    // Side view: Z,Y -> heatmap X,Y
                    x = Math.floor((pos.z / finalConfig.worldSize + 0.5) * finalConfig.width);
                    y = Math.floor((pos.y / finalConfig.worldSize + 0.5) * finalConfig.height);
                    break;
                case 'front':
                    // Front view: X,Y -> heatmap X,Y
                    x = Math.floor((pos.x / finalConfig.worldSize + 0.5) * finalConfig.width);
                    y = Math.floor((pos.y / finalConfig.worldSize + 0.5) * finalConfig.height);
                    break;
                default:
                    throw new Error(`Unknown view mode: ${finalConfig.viewMode}`);
            }
            
            if (x >= 0 && x < finalConfig.width && y >= 0 && y < finalConfig.height) {
                // Apply blur effect for smoother heatmap areas
                const blurRadius = finalConfig.blurAmount;
                for (let dx = -blurRadius; dx <= blurRadius; dx++) {
                    for (let dy = -blurRadius; dy <= blurRadius; dy++) {
                        const nx = x + dx;
                        const ny = y + dy;
                        
                        if (nx >= 0 && nx < finalConfig.width && ny >= 0 && ny < finalConfig.height) {
                            const distance = Math.sqrt(dx * dx + dy * dy);
                            // Gaussian-like falloff
                            const weight = Math.exp(-distance * distance / (2 * blurRadius * blurRadius));
                            heatmap[ny][nx] += Math.round(weight * 5);
                        }
                    }
                }
                good++;
            } else {
                bad++;
            }
        }
    }
    
    // console.log(`Time ${maxTime}s - Good: ${good}, Bad: ${bad}`);
    
    // Find maximum count for normalization
    let maxCount = 0;
    for (let y = 0; y < finalConfig.height; y++) {
        for (let x = 0; x < finalConfig.width; x++) {
            maxCount = Math.max(maxCount, heatmap[y][x]);
        }
    }
    
    // Create image data
    const ctx = canvas.getContext("2d");
    const imageData = ctx.createImageData(finalConfig.width, finalConfig.height);
    
    // Generate heatmap colors with transparency
    for (let y = 0; y < finalConfig.height; y++) {
        for (let x = 0; x < finalConfig.width; x++) {
            let value = maxCount > 0 ? heatmap[y][x] / maxCount : 0;
            
            // Apply power curve to increase sensitivity to lower values
            value = Math.pow(value, finalConfig.sensitivity);
            
            const color = getHeatmapColor(value);
            const index = (y * finalConfig.width + x) * 4;
            
            imageData.data[index] = color.r;     // Red
            imageData.data[index + 1] = color.g; // Green
            imageData.data[index + 2] = color.b; // Blue
            // Make heatmap semi-transparent based on intensity
            imageData.data[index + 3] = value > 0 ? Math.round(value * 200 + 55) : 0; // Alpha (0-255)
        }
    }
    
    // Scale the heatmap to a temporary canvas
    const tempCanvas = canvas.constructor as any;
    const tempCtx = new tempCanvas(finalConfig.width, finalConfig.height).getContext("2d");
    tempCtx.putImageData(imageData, 0, 0);
    
    // Clear main canvas
    ctx.fillStyle = "white";
    ctx.fillRect(0, 0, canvas.width, canvas.height);
    
    // Draw background image if provided (use preloaded image if available)
    if (finalConfig.preloadedBackground) {
        ctx.drawImage(finalConfig.preloadedBackground, 0, 0, canvas.width, canvas.height);
    } else if (finalConfig.backgroundImage) {
        try {
            const bgImage = await loadImage(finalConfig.backgroundImage);
            ctx.drawImage(bgImage, 0, 0, canvas.width, canvas.height);
        } catch (error) {
            console.error(`Failed to load background image: ${finalConfig.backgroundImage}`, error);
        }
    }
    
    // Draw scaled heatmap on top with transparency
    ctx.drawImage(tempCtx.canvas, 0, 0, canvas.width, canvas.height);
};

export const getReplayDuration = (replayData: ReplayData): number => {
    let maxTime = 0;
    
    for (const player of replayData.players) {
        for (const snapshot of player.snapshots) {
            maxTime = Math.max(maxTime, snapshot.time);
        }
    }
    
    return Math.ceil(maxTime); // Round up to get total seconds
};