src/components/performance/performance-comparison.tsx

"use client"

import { useState, useEffect } from "react"
import { useUser } from "@stackframe/stack"
import { Card, CardContent, CardHeader, CardTitle, CardDescription } from "@/components/ui/card"
import { Button } from "@/components/ui/button"
import { Badge } from "@/components/ui/badge"
import { Textarea } from "@/components/ui/textarea"
import { Input } from "@/components/ui/input"
import { Progress } from "@/components/ui/progress"
import { ShareMenu } from "@/components/share/share-menu"
import { useOllama } from "@/contexts/ollama-context"
import {
    BarChart3,
    Zap,
    Clock,
    DollarSign,
    Play,
    Loader2,
    Trophy,
    TrendingUp,
    TrendingDown,
    Save,
    Check,
    History,
    ChevronDown,
    ChevronUp,
    Rocket,
    Coins,
    Star,
    Server
} from "lucide-react"
import { ModelId } from "@/types/prompt"

interface ModelPerformance {
    model: string
    name: string
    provider: string
    responseTime: number
    estimatedTokens: number
    estimatedCost: number
    qualityScore: number
    output: string
    isOllama?: boolean
}

const MODEL_PRICING: Record<string, { input: number; output: number }> = {
    "gemini-2.0-flash": { input: 0.00001, output: 0.00004 },
    "gpt-4o": { input: 0.0025, output: 0.01 },
    "gpt-4o-mini": { input: 0.00015, output: 0.0006 },
    "claude-3-5-sonnet-latest": { input: 0.003, output: 0.015 },
}

const MODELS_TO_COMPARE: { id: ModelId; name: string; provider: string }[] = [
    { id: "gemini-2.0-flash", name: "Gemini 2.0 Flash", provider: "Google" },
    { id: "gpt-4o-mini", name: "GPT-4o Mini", provider: "OpenAI" },
    { id: "gpt-4o", name: "GPT-4o", provider: "OpenAI" },
    { id: "claude-3-5-sonnet", name: "Claude 3.5 Sonnet", provider: "Anthropic" },
]

interface SavedBenchmark {
    id: string
    name: string | null
    prompt: string
    results: { model: string; name: string; responseTime: number; tokens: number; cost: number }[]
    createdAt: string
}

export function PerformanceComparison() {
    const user = useUser()
    const { settings: ollamaSettings } = useOllama()
    const [prompt, setPrompt] = useState("")
    const [results, setResults] = useState<ModelPerformance[]>([])
    const [isRunning, setIsRunning] = useState(false)
    const [currentModel, setCurrentModel] = useState<string | null>(null)

    // Save state
    const [saving, setSaving] = useState(false)
    const [saved, setSaved] = useState(false)
    const [saveName, setSaveName] = useState("")
    const [savedBenchmarks, setSavedBenchmarks] = useState<SavedBenchmark[]>([])
    const [loadingHistory, setLoadingHistory] = useState(false)
    const [showHistory, setShowHistory] = useState(false)
    const [expandedId, setExpandedId] = useState<string | null>(null)

    // Fetch saved benchmarks
    useEffect(() => {
        const fetchHistory = async () => {
            if (!user?.id) return
            setLoadingHistory(true)
            try {
                const res = await fetch(`/api/comparisons?userId=${user.id}&type=performance`)
                if (res.ok) {
                    const data = await res.json()
                    setSavedBenchmarks(data.comparisons || [])
                }
            } catch (err) {
                console.error("Failed to fetch history:", err)
            } finally {
                setLoadingHistory(false)
            }
        }
        fetchHistory()
    }, [user?.id])

    const runComparison = async () => {
        if (!prompt.trim()) return

        setIsRunning(true)
        setResults([])

        const newResults: ModelPerformance[] = []

        // Combine cloud models with Ollama models
        const allModels = [
            ...MODELS_TO_COMPARE,
            ...ollamaSettings.availableModels.map(name => ({
                id: name,
                name: name,
                provider: "Ollama",
                isOllama: true
            }))
        ]

        for (const model of allModels) {
            setCurrentModel(model.name)
            const startTime = Date.now()

            try {
                const isOllama = 'isOllama' in model && model.isOllama
                let fullOutput = ""

                if (isOllama) {
                    // Ollama: stream directly from browser
                    const ollamaUrl = (ollamaSettings.apiUrl || "http://localhost:11434").replace(/\/+$/, "")
                    const response = await fetch(`${ollamaUrl}/api/generate`, {
                        method: "POST",
                        headers: { "Content-Type": "application/json" },
                        body: JSON.stringify({ model: model.id, prompt, stream: true }),
                    })

                    const reader = response.body?.getReader()
                    const decoder = new TextDecoder()

                    if (reader) {
                        while (true) {
                            const { done, value } = await reader.read()
                            if (done) break
                            const chunk = decoder.decode(value, { stream: true })
                            const lines = chunk.split("\n").filter(l => l.trim())
                            for (const line of lines) {
                                try {
                                    const json = JSON.parse(line)
                                    if (json.response) fullOutput += json.response
                                } catch { /* skip */ }
                            }
                        }
                    }
                } else {
                    // Cloud models: stream via server /api/run
                    const response = await fetch("/api/run", {
                        method: "POST",
                        headers: { "Content-Type": "application/json" },
                        body: JSON.stringify({ prompt, model: model.id }),
                    })

                    const reader = response.body?.getReader()
                    const decoder = new TextDecoder()

                    if (reader) {
                        while (true) {
                            const { done, value } = await reader.read()
                            if (done) break
                            fullOutput += decoder.decode(value)
                        }
                    }
                }

                const endTime = Date.now()
                const responseTime = endTime - startTime

                // Estimate tokens (rough: ~4 chars per token)
                const inputTokens = Math.ceil(prompt.length / 4)
                const outputTokens = Math.ceil(fullOutput.length / 4)
                const totalTokens = inputTokens + outputTokens

                // Calculate cost (Ollama is free)
                const pricing = isOllama ? { input: 0, output: 0 } : (MODEL_PRICING[model.id] || { input: 0, output: 0 })
                const cost = (inputTokens * pricing.input / 1000) + (outputTokens * pricing.output / 1000)

                newResults.push({
                    model: model.id,
                    name: model.name,
                    provider: model.provider,
                    responseTime,
                    estimatedTokens: totalTokens,
                    estimatedCost: cost,
                    qualityScore: 0,
                    output: fullOutput,
                    isOllama: isOllama,
                })

                setResults([...newResults])
            } catch (error) {
                console.error(`Error with ${model.name}:`, error)
            }
        }

        setCurrentModel(null)
        setIsRunning(false)
    }

    // Rate quality of a model's output
    const rateQuality = (modelId: string, score: number) => {
        setResults(prev => prev.map(r =>
            r.model === modelId ? { ...r, qualityScore: score } : r
        ))
    }

    // Find best performers
    const fastestModel = results.length > 0
        ? results.reduce((a, b) => a.responseTime < b.responseTime ? a : b)
        : null
    const cheapestModel = results.length > 0
        ? results.reduce((a, b) => a.estimatedCost < b.estimatedCost ? a : b)
        : null
    const mostEfficient = results.length > 0
        ? results.reduce((a, b) => (a.estimatedTokens / a.responseTime) > (b.estimatedTokens / b.responseTime) ? a : b)
        : null

    // Save benchmark
    const handleSave = async () => {
        if (results.length === 0) return
        setSaving(true)
        try {
            const res = await fetch('/api/comparisons', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({
                    userId: user?.id,
                    name: saveName || null,
                    type: 'performance',
                    prompt,
                    results: results.map(r => ({
                        model: r.model,
                        name: r.name,
                        responseTime: r.responseTime,
                        tokens: r.estimatedTokens,
                        cost: r.estimatedCost
                    })),
                    winner: fastestModel?.model
                })
            })
            if (res.ok) {
                const newBench = await res.json()
                setSavedBenchmarks(prev => [newBench, ...prev])
                setSaved(true)
                setSaveName("")
                setTimeout(() => setSaved(false), 3000)
            }
        } catch (err) {
            console.error("Failed to save:", err)
        } finally {
            setSaving(false)
        }
    }

    return (
        <div className="container mx-auto px-4 py-8 max-w-6xl">
            {/* Hero */}
            <div className="text-center mb-12">
                <div className="inline-flex items-center justify-center h-16 w-16 rounded-2xl bg-linear-to-br from-primary to-accent mb-4">
                    <BarChart3 className="h-8 w-8 text-white" />
                </div>
                <h1 className="text-4xl md:text-5xl font-serif font-medium mb-4">
                    Performance <span className="text-gradient italic">Comparison</span>
                </h1>
                <p className="text-xl text-muted-foreground max-w-2xl mx-auto">
                    Compare speed, token usage, and cost across AI models
                </p>
            </div>

            {/* Input */}
            <Card className="mb-8">
                <CardHeader>
                    <CardTitle className="text-lg">Test Prompt</CardTitle>
                    <CardDescription>Enter a prompt to benchmark across all models</CardDescription>
                </CardHeader>
                <CardContent className="space-y-4">
                    <Textarea
                        placeholder="Enter a prompt to test performance..."
                        value={prompt}
                        onChange={(e) => setPrompt(e.target.value)}
                        rows={4}
                    />
                    <div className="flex gap-3">
                        <Button
                            onClick={runComparison}
                            disabled={!prompt.trim() || isRunning}
                            className="gap-2"
                        >
                            {isRunning ? (
                                <>
                                    <Loader2 className="h-4 w-4 animate-spin" />
                                    Testing {currentModel}...
                                </>
                            ) : (
                                <>
                                    <Play className="h-4 w-4" />
                                    Run Benchmark
                                </>
                            )}
                        </Button>
                        {results.length > 0 && (
                            <ShareMenu
                                title="AI Performance Benchmark"
                                description={`Comparing ${MODELS_TO_COMPARE.length} AI models for speed and cost`}
                            />
                        )}
                    </div>
                </CardContent>
            </Card>

            {/* Winners */}
            {results.length === MODELS_TO_COMPARE.length && (
                <div className="grid md:grid-cols-3 gap-4 mb-8">
                    <Card className="border-green-500/50 bg-green-500/5">
                        <CardContent className="p-4 flex items-center gap-3">
                            <Clock className="h-8 w-8 text-green-500" />
                            <div>
                                <p className="text-xs text-muted-foreground">Fastest</p>
                                <p className="font-semibold">{fastestModel?.name}</p>
                                <p className="text-sm text-green-500">{(fastestModel?.responseTime || 0) / 1000}s</p>
                            </div>
                        </CardContent>
                    </Card>
                    <Card className="border-blue-500/50 bg-blue-500/5">
                        <CardContent className="p-4 flex items-center gap-3">
                            <DollarSign className="h-8 w-8 text-blue-500" />
                            <div>
                                <p className="text-xs text-muted-foreground">Cheapest</p>
                                <p className="font-semibold">{cheapestModel?.name}</p>
                                <p className="text-sm text-blue-500">${cheapestModel?.estimatedCost.toFixed(6)}</p>
                            </div>
                        </CardContent>
                    </Card>
                    <Card className="border-purple-500/50 bg-purple-500/5">
                        <CardContent className="p-4 flex items-center gap-3">
                            <Zap className="h-8 w-8 text-purple-500" />
                            <div>
                                <p className="text-xs text-muted-foreground">Most Efficient</p>
                                <p className="font-semibold">{mostEfficient?.name}</p>
                                <p className="text-sm text-purple-500">Best tokens/sec ratio</p>
                            </div>
                        </CardContent>
                    </Card>
                </div>
            )}

            {/* Results Table */}
            {results.length > 0 && (
                <Card>
                    <CardHeader>
                        <CardTitle className="text-lg">Benchmark Results</CardTitle>
                        <CardDescription>Comparison of {results.length} models</CardDescription>
                    </CardHeader>
                    <CardContent>
                        <div className="space-y-4">
                            {results.map((result, index) => {
                                const maxTime = Math.max(...results.map(r => r.responseTime))
                                const maxTokens = Math.max(...results.map(r => r.estimatedTokens))
                                const maxCost = Math.max(...results.map(r => r.estimatedCost))

                                return (
                                    <div key={result.model} className="p-4 rounded-lg border">
                                        <div className="flex items-center justify-between mb-4">
                                            <div className="flex items-center gap-2">
                                                <span className="text-2xl font-bold text-muted-foreground">
                                                    #{index + 1}
                                                </span>
                                                <div>
                                                    <p className="font-semibold">{result.name}</p>
                                                    <Badge variant="outline">{result.provider}</Badge>
                                                </div>
                                            </div>
                                            <div className="flex gap-2">
                                                {result === fastestModel && (
                                                    <Badge className="bg-green-500 flex items-center gap-1"><Rocket className="h-3 w-3" /> Fastest</Badge>
                                                )}
                                                {result === cheapestModel && (
                                                    <Badge className="bg-blue-500 flex items-center gap-1"><Coins className="h-3 w-3" /> Cheapest</Badge>
                                                )}
                                            </div>
                                        </div>

                                        <div className="grid md:grid-cols-3 gap-4">
                                            {/* Response Time */}
                                            <div>
                                                <div className="flex items-center justify-between text-sm mb-1">
                                                    <span className="flex items-center gap-1">
                                                        <Clock className="h-3 w-3" />
                                                        Response Time
                                                    </span>
                                                    <span>{(result.responseTime / 1000).toFixed(2)}s</span>
                                                </div>
                                                <Progress
                                                    value={(result.responseTime / maxTime) * 100}
                                                    className="h-2"
                                                />
                                            </div>

                                            {/* Tokens */}
                                            <div>
                                                <div className="flex items-center justify-between text-sm mb-1">
                                                    <span className="flex items-center gap-1">
                                                        <Zap className="h-3 w-3" />
                                                        Est. Tokens
                                                    </span>
                                                    <span>{result.estimatedTokens.toLocaleString()}</span>
                                                </div>
                                                <Progress
                                                    value={(result.estimatedTokens / maxTokens) * 100}
                                                    className="h-2"
                                                />
                                            </div>

                                            {/* Cost */}
                                            <div>
                                                <div className="flex items-center justify-between text-sm mb-1">
                                                    <span className="flex items-center gap-1">
                                                        <DollarSign className="h-3 w-3" />
                                                        Est. Cost
                                                    </span>
                                                    <span className="flex items-center gap-1">
                                                        {result.isOllama && <Badge variant="outline" className="text-xs"><Server className="h-2 w-2 mr-1" />Local</Badge>}
                                                        ${result.estimatedCost.toFixed(6)}
                                                    </span>
                                                </div>
                                                <Progress
                                                    value={maxCost > 0 ? (result.estimatedCost / maxCost) * 100 : 0}
                                                    className="h-2"
                                                />
                                            </div>
                                        </div>

                                        {/* Quality Rating */}
                                        <div className="mt-4 pt-4 border-t flex items-center justify-between">
                                            <div className="flex items-center gap-2">
                                                <span className="text-sm text-muted-foreground">Rate Quality:</span>
                                                <div className="flex gap-1">
                                                    {[1, 2, 3, 4, 5].map(score => (
                                                        <button
                                                            key={score}
                                                            onClick={() => rateQuality(result.model, score)}
                                                            className="p-1 hover:scale-110 transition-transform"
                                                        >
                                                            <Star
                                                                className={`h-5 w-5 ${score <= result.qualityScore ? 'fill-yellow-400 text-yellow-400' : 'text-muted-foreground'}`}
                                                            />
                                                        </button>
                                                    ))}
                                                </div>
                                                {result.qualityScore > 0 && (
                                                    <span className="text-sm font-medium">{result.qualityScore}/5</span>
                                                )}
                                            </div>
                                            {result.output && (
                                                <Badge variant="secondary" className="text-xs">
                                                    {result.output.length} chars
                                                </Badge>
                                            )}
                                        </div>
                                    </div>
                                )
                            })}
                        </div>
                    </CardContent>
                </Card>
            )}

            {/* Save Benchmark */}
            {results.length === MODELS_TO_COMPARE.length && (
                <Card className="mt-6 border-primary/50">
                    <CardHeader className="pb-3">
                        <div className="flex items-center gap-2">
                            <Save className="h-5 w-5 text-primary" />
                            <CardTitle className="text-base">Save Benchmark</CardTitle>
                        </div>
                        <CardDescription>Save these results to compare later</CardDescription>
                    </CardHeader>
                    <CardContent>
                        <div className="flex flex-col sm:flex-row gap-3">
                            <div className="flex-1">
                                <Input
                                    value={saveName}
                                    onChange={(e) => setSaveName(e.target.value)}
                                    placeholder="Name this benchmark (optional)"
                                />
                            </div>
                            <Button
                                onClick={handleSave}
                                disabled={saving || saved}
                                className="gap-2"
                            >
                                {saving ? (
                                    <><Loader2 className="h-4 w-4 animate-spin" />Saving...</>
                                ) : saved ? (
                                    <><Check className="h-4 w-4" />Saved!</>
                                ) : (
                                    <><Save className="h-4 w-4" />Save Result</>
                                )}
                            </Button>
                        </div>
                    </CardContent>
                </Card>
            )}

            {/* Saved Benchmarks History */}
            <Card className="mt-6">
                <CardHeader className="pb-3">
                    <button
                        onClick={() => setShowHistory(!showHistory)}
                        className="flex items-center justify-between w-full text-left"
                    >
                        <div className="flex items-center gap-2">
                            <History className="h-5 w-5 text-primary" />
                            <CardTitle className="text-base">Saved Benchmarks</CardTitle>
                            <Badge variant="outline">{savedBenchmarks.length}</Badge>
                        </div>
                        {showHistory ? (
                            <ChevronUp className="h-5 w-5 text-muted-foreground" />
                        ) : (
                            <ChevronDown className="h-5 w-5 text-muted-foreground" />
                        )}
                    </button>
                    <CardDescription>View previously saved benchmarks</CardDescription>
                </CardHeader>
                {showHistory && (
                    <CardContent>
                        {loadingHistory ? (
                            <div className="flex items-center justify-center py-8">
                                <Loader2 className="h-6 w-6 animate-spin text-muted-foreground" />
                            </div>
                        ) : savedBenchmarks.length === 0 ? (
                            <p className="text-muted-foreground text-center py-8">
                                No saved benchmarks yet. Run a benchmark and save it!
                            </p>
                        ) : (
                            <div className="space-y-3">
                                {savedBenchmarks.map((bench) => (
                                    <div key={bench.id} className="border rounded-lg">
                                        <button
                                            onClick={() => setExpandedId(expandedId === bench.id ? null : bench.id)}
                                            className="w-full p-3 flex items-center justify-between text-left hover:bg-muted/50"
                                        >
                                            <div>
                                                <p className="font-medium text-sm">
                                                    {bench.name || `Benchmark from ${new Date(bench.createdAt).toLocaleDateString()}`}
                                                </p>
                                                <p className="text-xs text-muted-foreground">
                                                    {new Date(bench.createdAt).toLocaleString()}
                                                </p>
                                            </div>
                                            {expandedId === bench.id ? (
                                                <ChevronUp className="h-4 w-4 text-muted-foreground" />
                                            ) : (
                                                <ChevronDown className="h-4 w-4 text-muted-foreground" />
                                            )}
                                        </button>
                                        {expandedId === bench.id && (
                                            <div className="p-3 border-t space-y-2">
                                                <p className="text-sm text-muted-foreground mb-2">Prompt: {bench.prompt.slice(0, 100)}{bench.prompt.length > 100 ? '...' : ''}</p>
                                                {bench.results.map((r, idx) => (
                                                    <div key={idx} className="flex items-center justify-between text-sm border-b pb-1">
                                                        <span className="font-medium">{r.name}</span>
                                                        <div className="flex gap-4 text-muted-foreground">
                                                            <span>{(r.responseTime / 1000).toFixed(2)}s</span>
                                                            <span>{r.tokens} tokens</span>
                                                            <span>${r.cost.toFixed(6)}</span>
                                                        </div>
                                                    </div>
                                                ))}
                                            </div>
                                        )}
                                    </div>
                                ))}
                            </div>
                        )}
                    </CardContent>
                )}
            </Card>
        </div>
    )
}