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

const DEFAULT_MODEL_URL = "https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf";

export const ChatPage = () => {
    const [messages, setMessages] = useState<ChatMessage[]>([])
    const [input, setInput] = useState('')
    const [isTyping, setIsTyping] = useState(false)
    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)
    const scrollRef = useRef<HTMLDivElement>(null)
    const [isServerReady, setIsServerReady] = useState(false);

    // Download settings
    const [modelUrl, setModelUrl] = useState(DEFAULT_MODEL_URL);
    const [authToken, setAuthToken] = useState('');
    const [showAdvanced, setShowAdvanced] = useState(false);

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

    const checkModel = async () => {
        try {
            const exists = await LocalAIService.checkModelExists();
            if (exists) {
                setModelStatus('ready');
                initServer();
            } else {
                setModelStatus('missing');
            }
        } catch (e) {
            console.error(e);
            setErrorMsg("Failed to check model status.");
            setModelStatus('error');
        }
    }

    const initServer = async () => {
        try {
            await LocalAIService.startServer();
            setIsServerReady(true);
        } catch (e) {
            // It might already be running, or failed. 
            // We'll assume if it fails it might be because it's already on.
            // A robust health check endpoint would be better.
            console.warn("Server start response:", e);
            setIsServerReady(true); // optimistically assume ready or already running
        }
    }

    const handleDownload = async () => {
        setModelStatus('downloading');
        setErrorMsg(null);
        try {
            await LocalAIService.downloadModel(modelUrl, authToken || undefined, (current, total) => {
                setDownloadProgress({ current, total });
            });
            setModelStatus('ready');
            initServer();
        } catch (e) {
            console.error(e);
            setErrorMsg(`Failed to download model. ${e}`);
            setModelStatus('missing'); // Go back to missing so they can try again
            setDownloadProgress(null);
        }
    }

    useEffect(() => {
        if (scrollRef.current) {
            scrollRef.current.scrollTop = scrollRef.current.scrollHeight
        }
    }, [messages, isTyping])

    const handleSend = async () => {
        if (!input.trim() || !isServerReady) return

        const userMsg: ChatMessage = {
            id: Date.now().toString(),
            role: 'user',
            content: input,
            timestamp: new Date()
        }

        setMessages(prev => [...prev, userMsg])
        setInput('')
        setIsTyping(true)

        try {
            // Convert history to format expected by AI
            // We might want to keep a 'system' prompt first
            const apiMessages = [
                { role: "system", content: "You are a helpful Norwegian language tutor. We are having a casual conversation in Norwegian. Please reply in Norwegian." },
                ...messages.map(m => ({ role: m.role, content: m.content })),
                { role: "user", content: userMsg.content }
            ];

            const response = await LocalAIService.chatCompletion(apiMessages);

            const aiMsg: ChatMessage = {
                id: (Date.now() + 1).toString(),
                role: 'assistant',
                content: response,
                timestamp: new Date()
            }
            setMessages(prev => [...prev, aiMsg])
        } catch (e) {
            console.error(e);
            const errorMsg: ChatMessage = {
                id: Date.now().toString(),
                role: 'assistant',
                content: "Beklager, jeg fikk problemer med å koble til AI-tjenesten.",
                timestamp: new Date()
            }
            setMessages(prev => [...prev, errorMsg])
        } finally {
            setIsTyping(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">
                    <Download className="w-10 h-10" />
                </div>
                <h2 className="text-2xl font-bold">AI Model Missing</h2>
                <p className="text-muted-foreground">
                    To use the AI tutor details, we need to download a lightweight language model (approx 2.5GB).
                    This is a one-time setup.
                </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>
                ) : (
                    <div className="w-full space-y-4">
                        <Button onClick={handleDownload} size="lg" className="w-full sm:w-auto">
                            Download Model
                        </Button>

                        <div className="pt-4 border-t w-full">
                            <button
                                onClick={() => setShowAdvanced(!showAdvanced)}
                                className="text-xs text-muted-foreground hover:underline"
                            >
                                {showAdvanced ? "Hide Advanced Options" : "Advanced Options (Custom URL / Token)"}
                            </button>

                            {showAdvanced && (
                                <div className="space-y-3 mt-3 text-left">
                                    <div className="space-y-1">
                                        <label className="text-xs font-medium">Model URL (GGUF)</label>
                                        <Input
                                            value={modelUrl}
                                            onChange={(e) => setModelUrl(e.target.value)}
                                            className="text-xs"
                                        />
                                    </div>
                                    <div className="space-y-1">
                                        <label className="text-xs font-medium">HuggingFace Token (Optional)</label>
                                        <Input
                                            type="password"
                                            value={authToken}
                                            onChange={(e) => setAuthToken(e.target.value)}
                                            placeholder="hf_..."
                                            className="text-xs"
                                        />
                                        <p className="text-[10px] text-muted-foreground">Required only for gated models.</p>
                                    </div>
                                </div>
                            )}
                        </div>
                    </div>
                )}
            </div>
        )
    }

    if (modelStatus === 'error') {
        return (
            <div className="flex flex-col h-full items-center justify-center p-8 text-center">
                <AlertCircle className="w-16 h-16 text-destructive mb-4" />
                <h2 className="text-xl font-bold text-destructive">Something went wrong</h2>
                <p className="text-muted-foreground mt-2">{errorMsg}</p>
                <Button onClick={checkModel} variant="outline" className="mt-6">Retry</Button>
            </div>
        )
    }

    return (
        <div className="flex flex-col h-[calc(100vh-2rem)] max-w-4xl mx-auto p-4 md:p-6">
            <header className="flex items-center justify-between mb-6 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">
                        <Bot className="w-6 h-6" />
                    </div>
                    <div>
                        <h1 className="text-xl font-bold">Lærer AI</h1>
                        <p className="text-xs text-muted-foreground flex items-center gap-1">
                            <span className={cn("w-2 h-2 rounded-full", isServerReady ? "bg-green-500 animate-pulse" : "bg-yellow-500")} />
                            {isServerReady ? "Online" : "Initializing..."}
                        </p>
                    </div>
                </div>
                <Button variant="ghost" size="icon" onClick={() => setMessages([])}>
                    <RefreshCw className="w-5 h-5 text-muted-foreground" />
                </Button>
            </header>

            <div
                ref={scrollRef}
                className="flex-1 overflow-y-auto space-y-6 mb-6 px-2 scroll-smooth"
            >
                {messages.length === 0 && (
                    <div className="h-full flex flex-col items-center justify-center text-muted-foreground opacity-50">
                        <Bot className="w-16 h-16 mb-4" />
                        <p>Start a conversation to practice your Norwegian!</p>
                    </div>
                )}

                {messages.map((msg) => (
                    <div
                        key={msg.id}
                        className={cn(
                            "flex gap-4 max-w-[80%]",
                            msg.role === 'user' ? "ml-auto flex-row-reverse" : ""
                        )}
                    >
                        <div className={cn(
                            "w-8 h-8 rounded-full flex items-center justify-center shrink-0 mt-1",
                            msg.role === 'user' ? "bg-primary text-primary-foreground" : "bg-muted text-muted-foreground"
                        )}>
                            {msg.role === 'user' ? <User className="w-5 h-5" /> : <Bot className="w-5 h-5" />}
                        </div>

                        <div className={cn(
                            "p-4 rounded-2xl text-sm leading-relaxed shadow-sm",
                            msg.role === 'user'
                                ? "bg-primary text-primary-foreground rounded-tr-sm"
                                : "bg-card border rounded-tl-sm"
                        )}>
                            {msg.content}
                            <span className="block text-[10px] opacity-50 mt-2">
                                {msg.timestamp.toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' })}
                            </span>
                        </div>
                    </div>
                ))}

                {isTyping && (
                    <div className="flex gap-4 max-w-[80%]">
                        <div className="w-8 h-8 rounded-full bg-muted flex items-center justify-center shrink-0 mt-1">
                            <Bot className="w-5 h-5 text-muted-foreground" />
                        </div>
                        <div className="bg-card border p-4 rounded-2xl rounded-tl-sm shadow-sm flex items-center gap-1">
                            <span className="w-1.5 h-1.5 bg-muted-foreground rounded-full animate-bounce [animation-delay:-0.3s]" />
                            <span className="w-1.5 h-1.5 bg-muted-foreground rounded-full animate-bounce [animation-delay:-0.15s]" />
                            <span className="w-1.5 h-1.5 bg-muted-foreground rounded-full animate-bounce" />
                        </div>
                    </div>
                )}
            </div>

            <div className="flex gap-2 bg-card p-2 rounded-lg border shadow-sm">
                <Input
                    placeholder="Type your message in Norwegian..."
                    value={input}
                    onChange={(e) => setInput(e.target.value)}
                    onKeyDown={(e) => e.key === 'Enter' && handleSend()}
                    disabled={!isServerReady}
                    className="border-0 focus-visible:ring-0 bg-transparent text-base"
                />
                <Button onClick={handleSend} size="icon" disabled={!isServerReady} className={cn("transition-all", input.trim() ? "scale-100" : "scale-90 opacity-50")}>
                    <Send className="w-5 h-5" />
                </Button>
            </div>
        </div>
    )
}
