import { useState, useRef, useEffect } from 'react'
import { Button } from '@/components/ui/button'
import { Input } from '@/components/ui/input'
import { Send, Bot, RefreshCw, Download, Loader2, AlertCircle, Image as ImageIcon, Sparkles } from 'lucide-react'
import { cn } from '@/lib/utils'
import { LocalAIService } from '@/services/localAi'

const VISION_MODEL_URL = "https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf/resolve/main/ggml-model-Q4_K_M.gguf";
const VISION_PROJ_URL = "https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf/resolve/main/mmproj-model-f16.gguf";

export const VisualRecognitionPage = () => {
    const [modelStatus, setModelStatus] = useState<'checking' | 'missing' | 'downloading' | 'ready' | 'error'>('checking')
    const [downloadProgress, setDownloadProgress] = useState<{ current: number; total: number } | null>(null)
    const [errorMsg, setErrorMsg] = useState<string | null>(null)

    // Image State
    const [imageSeed, setImageSeed] = useState(Date.now());
    const [imageUrl, setImageUrl] = useState(`https://picsum.photos/seed/${Date.now()}/600/400`);
    const [userDescription, setUserDescription] = useState('');
    const [aiDescription, setAiDescription] = useState<string | null>(null);
    const [feedback, setFeedback] = useState<string | null>(null);
    const [isAnalyzing, setIsAnalyzing] = useState(false);
    const [isGeneratingFeedback, setIsGeneratingFeedback] = useState(false);

    useEffect(() => {
        checkVisionModel();
    }, [])

    const checkVisionModel = async () => {
        try {
            const exists = await LocalAIService.checkVisionModelExists();
            setModelStatus(exists ? 'ready' : 'missing');
        } catch (e) {
            console.error(e);
            setModelStatus('error');
            setErrorMsg("Failed to check vision model.");
        }
    }

    const handleDownload = async () => {
        setModelStatus('downloading');
        setErrorMsg(null);
        try {
            await LocalAIService.downloadVisionModel(VISION_MODEL_URL, VISION_PROJ_URL, undefined, (current, total) => {
                setDownloadProgress({ current, total });
            });
            setModelStatus('ready');
        } catch (e) {
            console.error(e);
            setErrorMsg(`Failed to download vision model. ${e}`);
            setModelStatus('missing');
            setDownloadProgress(null);
        }
    }

    const handleRefreshImage = () => {
        const seed = Date.now();
        setImageSeed(seed);
        setImageUrl(`https://picsum.photos/seed/${seed}/600/400`);
        setUserDescription('');
        setAiDescription(null);
        setFeedback(null);
    }

    // Helper to ensure chat server is running
    const ensureChatServerReady = async () => {
        try {
            // Optimistically try to start - if already running, backend handles it (or just spawns another? actually backend just spawns.)
            // Logic: Try to fetch first. If fail, start, then poll.

            const checkServer = async () => {
                try {
                    // Try a lightweight request (e.g. models or just connect)
                    // We don't have a dedicated health endpoint, checking completion endpoint with empty body or GET
                    // usually llama-server has /health. Let's try /health
                    const res = await fetch("http://localhost:8080/health");
                    return res.ok;
                } catch {
                    return false;
                }
            };

            if (await checkServer()) return true;

            // Start server
            console.log("Starting Chat Server...");
            await LocalAIService.startServer();

            // Poll for up to 30s
            for (let i = 0; i < 10; i++) {
                await new Promise(r => setTimeout(r, 3000)); // wait 3s
                if (await checkServer()) return true;
                console.log("Waiting for server...");
            }
            return false;
        } catch (e) {
            console.error("Error ensuring chat server:", e);
            return false;
        }
    }

    const handleAnalyze = async () => {
        if (!userDescription.trim()) return;

        setIsAnalyzing(true);
        setAiDescription(null);
        setFeedback(null);
        setErrorMsg(null);

        let analysisResult = "";

        try {
            // 1. Analyze Image
            const analysis = await LocalAIService.analyzeImage(imageUrl);
            setAiDescription(analysis);
            analysisResult = analysis;

        } catch (e) {
            console.error(e);
            setErrorMsg("Analysis failed. Please try again.");
            setIsAnalyzing(false);
            return;
        }

        // 2. Generate Feedback using Instruct Model
        try {
            setIsGeneratingFeedback(true);
            setIsAnalyzing(false); // Stop analyzing spinner, start feedback spinner? 
            // Actually button logic uses OR. So spinner continues. Correct.

            const serverReady = await ensureChatServerReady();
            if (!serverReady) {
                throw new Error("Chat Server failed to start. Please try later.");
            }

            const prompt = `
You are a vision feedback assistant. 
I will provide you with the TRUE description of an image (generated by a vision model) and a USER description of the same image.
Your goal is to compare them and give the user feedback on their observation skills or language accuracy.

TRUE DESCRIPTION:
${analysisResult}

USER DESCRIPTION:
${userDescription}

Please provide constructive feedback in Norwegian.
            `.trim();

            const messages = [
                { role: 'system', content: 'You are a helpful Norwegian tutor.' },
                { role: 'user', content: prompt }
            ];

            const response = await LocalAIService.chatCompletion(messages);
            setFeedback(response);

        } catch (e) {
            console.error(e);
            setErrorMsg(`Image Analyzed, but Feedback failed: ${e}`);
        } finally {
            setIsAnalyzing(false);
            setIsGeneratingFeedback(false);
        }
    }

    if (modelStatus === 'checking') {
        return (
            <div className="flex h-full items-center justify-center">
                <Loader2 className="w-8 h-8 animate-spin text-primary" />
            </div>
        )
    }

    if (modelStatus === 'missing' || modelStatus === 'downloading') {
        return (
            <div className="flex flex-col h-full items-center justify-center p-8 text-center max-w-lg mx-auto space-y-6">
                <div className="w-20 h-20 bg-primary/10 rounded-full flex items-center justify-center text-primary mb-4">
                    <ImageIcon className="w-10 h-10" />
                </div>
                <h2 className="text-2xl font-bold">Vision Model Required</h2>
                <p className="text-muted-foreground">
                    To use the Visual Recognition feature, we need to download the MiniCPM-V 2.6 model (approx 5GB).
                    This enables "sight" for the AI.
                </p>

                {errorMsg && (
                    <div className="text-destructive text-sm bg-destructive/10 p-2 rounded">
                        {errorMsg}
                    </div>
                )}

                {modelStatus === 'downloading' && downloadProgress ? (
                    <div className="w-full space-y-2">
                        <div className="h-2 bg-secondary rounded-full overflow-hidden">
                            <div
                                className="h-full bg-primary transition-all duration-300"
                                style={{ width: `${(downloadProgress.current / downloadProgress.total) * 100}%` }}
                            />
                        </div>
                        <p className="text-sm text-muted-foreground">
                            {(downloadProgress.current / 1024 / 1024).toFixed(1)} MB / {(downloadProgress.total / 1024 / 1024).toFixed(1)} MB
                        </p>
                    </div>
                ) : (
                    <Button onClick={handleDownload} size="lg" className="w-full sm:w-auto">
                        Download Vision Model
                    </Button>
                )}
            </div>
        )
    }

    return (
        <div className="flex flex-col h-[calc(100vh-2rem)] max-w-5xl mx-auto p-4 md:p-6 gap-6">
            <header className="flex items-center justify-between pb-4 border-b">
                <div className="flex items-center gap-4">
                    <div className="w-10 h-10 rounded-full bg-primary/10 flex items-center justify-center text-primary">
                        <ImageIcon className="w-6 h-6" />
                    </div>
                    <div>
                        <h1 className="text-xl font-bold">Visual Recognition</h1>
                        <p className="text-xs text-muted-foreground">Describe the image and get AI feedback</p>
                    </div>
                </div>
                <Button variant="ghost" size="icon" onClick={handleRefreshImage}>
                    <RefreshCw className="w-5 h-5 text-muted-foreground" />
                </Button>
            </header>

            <div className="grid grid-cols-1 md:grid-cols-2 gap-6 flex-1 overflow-hidden">
                {/* Left Column: Image & Input */}
                <div className="flex flex-col gap-4 overflow-y-auto pr-2">
                    <div className="relative aspect-video bg-muted rounded-xl overflow-hidden shadow-sm border">
                        <img
                            src={imageUrl}
                            alt="Random scene"
                            className="w-full h-full object-cover"
                            key={imageSeed}
                        />
                    </div>

                    <div className="space-y-2">
                        <label className="text-sm font-medium">Describe what you see (in Norwegian):</label>
                        <textarea
                            className="w-full h-32 p-3 text-sm rounded-md border bg-transparent resize-none focus:ring-1 focus:ring-primary outline-none"
                            placeholder="Example: Jeg ser et vakkert fjell..."
                            value={userDescription}
                            onChange={(e) => setUserDescription(e.target.value)}
                        />
                        <Button
                            onClick={handleAnalyze}
                            disabled={!userDescription.trim() || isAnalyzing || isGeneratingFeedback}
                            className="w-full"
                        >
                            {(isAnalyzing || isGeneratingFeedback) ? (
                                <>
                                    <Loader2 className="w-4 h-4 mr-2 animate-spin" />
                                    {isAnalyzing ? "Analyzing Image..." : "Generating Feedback..."}
                                </>
                            ) : (
                                <>
                                    <Sparkles className="w-4 h-4 mr-2" />
                                    Analyze & Compare
                                </>
                            )}
                        </Button>
                    </div>
                </div>

                {/* Right Column: Results */}
                <div className="flex flex-col gap-4 h-full bg-muted/30 rounded-xl p-4 border overflow-y-auto">
                    {!aiDescription && !feedback && (
                        <div className="flex flex-col items-center justify-center h-full text-muted-foreground opacity-50 text-center p-4">
                            <Sparkles className="w-12 h-12 mb-3" />
                            <p>Describe the image on the left and click Analyze to see how your vision compares to the AI's.</p>
                        </div>
                    )}

                    {aiDescription && (
                        <div className="space-y-2 animate-in fade-in slide-in-from-bottom-4 duration-500">
                            <h3 className="text-sm font-semibold flex items-center gap-2">
                                <Bot className="w-4 h-4 text-primary" />
                                AI's Observation (Hidden Truth)
                            </h3>
                            <div className="text-xs bg-card p-3 rounded-md border text-muted-foreground leading-relaxed">
                                {aiDescription}
                            </div>
                        </div>
                    )}

                    {feedback && (
                        <div className="space-y-2 pt-4 border-t animate-in fade-in slide-in-from-bottom-4 duration-700 delay-200">
                            <h3 className="text-sm font-semibold flex items-center gap-2 text-primary">
                                <Sparkles className="w-4 h-4" />
                                Feedback
                            </h3>
                            <div className="text-sm bg-primary/5 p-4 rounded-md border border-primary/10 leading-relaxed whitespace-pre-wrap">
                                {feedback}
                            </div>
                        </div>
                    )}
                </div>
            </div>
        </div>
    )
}
