// src/app/samples/page.tsx
"use client"

import {useSamples} from '@/contexts/SamplesContext'
import Link from "next/link"
import {useMemo} from 'react'
import {useTraining} from "@/contexts/TrainingContext";
import {SampleCard} from "@/components/SampleCard";

// Helper function to parse sample information
const parseSampleInfo = (filename: string) => {
    const [timestamp, info] = filename.split('__')
    const [batch, index] = info.split('.')[0].split('_')
    return {
        timestamp: parseInt(timestamp),
        batch: parseInt(batch),
        index: parseInt(index)
    }
}

export default function SamplesPage() {
    const {samples, isLoading: samplesLoading, error: samplesError} = useSamples()
    const {config, isLoading: configLoading} = useTraining()

    // Get prompts from config
    const prompts = config?.config?.process[0]?.sample?.prompts || []

    // Group samples by batch number using a memoized calculation
    const groupedSamples = useMemo(() => {
        if (!samples?.length) return new Map()

        // Create groups based on batch numbers
        const groups = samples.reduce((acc, sample) => {
            const {batch} = parseSampleInfo(sample.filename)
            const group = acc.get(batch) || []
            group.push({
                ...sample,
                ...parseSampleInfo(sample.filename)
            })
            acc.set(batch, group)
            return acc
        }, new Map())

        // Sort samples within each group by index
        for (const [batch, items] of groups) {
            groups.set(batch, items.sort((a: any, b: any) => a.index - b.index))
        }

        // Return groups sorted by batch number (descending)
        return new Map([...groups].sort((a, b) => b[0] - a[0]))
    }, [samples])

    // Handle loading and error states
    if (samplesLoading || configLoading) return (
        <div className="min-h-screen bg-gray-900 p-4">
            <div className="max-w-7xl mx-auto text-gray-400">Loading...</div>
        </div>
    )

    if (samplesError) return (
        <div className="min-h-screen bg-gray-900 p-4">
            <div className="max-w-7xl mx-auto text-red-400">Error: {samplesError.message}</div>
        </div>
    )


    return (
        <main className="min-h-screen bg-gray-900 p-4">
            <div className="max-w-7xl mx-auto space-y-6">
                {/* Header with navigation */}
                <header className="flex justify-between items-center">
                    <h1 className="text-3xl font-bold text-gray-100">Samples Gallery</h1>
                    <Link
                        href="/"
                        className="text-gray-400 hover:text-gray-200 transition-colors duration-200 flex items-center gap-2 group"
                    >
                        ← Back to Dashboard
                    </Link>
                </header>


                {/* Prompts display */}
                <div className="sticky top-4 z-10 bg-gray-800 rounded-lg shadow-lg border border-gray-700 p-6">
                    <h2 className="text-xl font-semibold text-gray-300 mb-4">
                        Generation Prompts
                    </h2>
                    <div className="grid grid-cols-[repeat(auto-fit,minmax(200px,1fr))] gap-4">
                        {prompts.map((prompt, index) => (
                            <div
                                key={index}
                                className="text-sm text-gray-400 p-3 bg-gray-700 rounded-lg"
                            >
                                <span className="font-mono text-gray-500 mr-2">{index}:</span>
                                {prompt}
                            </div>
                        ))}
                    </div>
                </div>

                {/* Sample groups */}
                <div className="space-y-8">
                    {Array.from(groupedSamples).map(([batch, items]) => (
                        <div key={batch} className="bg-gray-800 rounded-lg shadow-lg border border-gray-700 p-6">
                            <h2 className="text-xl font-semibold text-gray-300 mb-4">
                                Steps {batch}
                            </h2>

                            {/* Scrollable container for samples */}
                            <div className="overflow-x-auto">
                                <div className="flex gap-4 min-w-full pb-4">
                                    {items.map((sample: any) => (
                                        <SampleCard
                                            key={sample.filename}
                                            sample={sample}
                                            prompt={prompts[sample.index]}
                                        />
                                    ))}
                                </div>
                            </div>
                        </div>
                    ))}
                </div>
            </div>
        </main>
    )
}