import { useState, useRef, useEffect, useCallback } from "react";
import { Send, RotateCcw, Zap } from "lucide-react";
import { useLLM } from "./hooks/useLLM";
import { LoadingScreen } from "./components/LoadingScreen";
import { ChatMessage } from "./components/ChatMessage";

interface Message {
  role: "user" | "assistant";
  content: string;
}

const SYSTEM_MESSAGE = {
  role: "system",
  content:
    "You are Maincoder, an expert code generation assistant. Write clean, well-structured code. When responding with code, use markdown code blocks with the appropriate language identifier.",
};

const EXAMPLES = [
  {
    icon: "🐍",
    text: "Binary search in Python",
    message: "Write a binary search function in Python",
  },
  {
    icon: "⚡",
    text: "Fibonacci with memoization",
    message: "Write a fibonacci function with memoization in JavaScript",
  },
];

const App = () => {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState("");
  const [isGenerating, setIsGenerating] = useState(false);
  const chatRef = useRef<HTMLDivElement>(null);
  const inputRef = useRef<HTMLInputElement>(null);

  const { isLoading, isReady, error, progress, loadModel, generateResponse, clearHistory } =
    useLLM();

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

  const clearChat = useCallback(() => {
    setMessages([]);
    clearHistory();
  }, [clearHistory]);

  const sendMessage = useCallback(
    async (text: string) => {
      if (!text.trim() || !isReady || isGenerating) return;

      const userMessage: Message = { role: "user", content: text };
      const currentMessages = [...messages, userMessage];
      setMessages(currentMessages);
      setInput("");
      setIsGenerating(true);

      try {
        const messagesForModel = [
          SYSTEM_MESSAGE,
          ...currentMessages.map((m) => ({ role: m.role, content: m.content })),
        ];

        setMessages([...currentMessages, { role: "assistant", content: "" }]);

        let accumulated = "";
        const response = await generateResponse(
          messagesForModel,
          (token: string) => {
            accumulated += token;
            setMessages((prev) => {
              const updated = [...prev];
              updated[updated.length - 1] = {
                role: "assistant",
                content: accumulated,
              };
              return updated;
            });
          },
        );

        setMessages([
          ...currentMessages,
          { role: "assistant", content: response },
        ]);
      } catch (err) {
        const errorMsg =
          err instanceof Error ? err.message : "Generation failed";
        setMessages([
          ...currentMessages,
          { role: "assistant", content: `Error: ${errorMsg}` },
        ]);
      } finally {
        setIsGenerating(false);
        setTimeout(() => inputRef.current?.focus(), 0);
      }
    },
    [messages, isReady, isGenerating, generateResponse],
  );

  if (!isReady) {
    return (
      <LoadingScreen
        isLoading={isLoading}
        progress={progress}
        error={error}
        onLoad={loadModel}
      />
    );
  }

  return (
    <div className="font-sans min-h-screen bg-gradient-to-br from-[#0a0f1a] via-[#0d1520] to-[#060a12] text-gray-100">
      <div className="flex h-screen w-full py-6 px-4 md:py-10 md:px-8">
        <div className="flex-1 flex flex-col p-6 bg-white/5 backdrop-blur-lg border border-white/10 rounded-3xl shadow-[0_35px_65px_rgba(0,0,0,0.5)] min-h-0">
          {/* Header */}
          <div className="flex items-center justify-between mb-6">
            <div className="space-y-1">
              <div className="flex items-center gap-2">
                <Zap size={16} className="text-emerald-400" />
                <span className="text-xs font-semibold uppercase tracking-[0.35em] text-emerald-400">
                  WebGPU
                </span>
              </div>
              <h1 className="text-2xl md:text-3xl font-bold text-white">
                Maincoder <span className="text-emerald-400">1B</span>
              </h1>
            </div>
            <button
              disabled={isGenerating}
              onClick={clearChat}
              className={`h-10 flex items-center px-4 rounded-full font-semibold text-sm transition-all border ${
                isGenerating
                  ? "border-white/10 bg-white/5 text-gray-500 cursor-not-allowed"
                  : "border-white/15 bg-white/8 text-gray-300 hover:border-emerald-500/40 hover:bg-emerald-500/10"
              }`}
            >
              <RotateCcw size={14} className="mr-2" /> New Chat
            </button>
          </div>

          {/* Chat Area */}
          <div
            ref={chatRef}
            className="flex-grow bg-[#080d16]/80 border border-white/10 rounded-2xl p-6 overflow-y-auto mb-6 space-y-5 shadow-inner min-h-0"
          >
            {messages.length === 0 ? (
              <div className="flex flex-col items-center justify-center h-full space-y-6">
                <div className="text-center mb-4">
                  <h2 className="text-2xl font-semibold text-white mb-1">
                    What would you like to code?
                  </h2>
                  <p className="text-sm text-gray-400">
                    Try an example or type your own prompt
                  </p>
                </div>
                <div className="grid grid-cols-1 sm:grid-cols-2 gap-3 max-w-2xl w-full px-4">
                  {EXAMPLES.map((example, i) => (
                    <button
                      key={i}
                      onClick={() => sendMessage(example.message)}
                      className="group relative overflow-hidden rounded-2xl border border-white/10 bg-white/5 text-left transition-all hover:-translate-y-0.5 hover:shadow-[0_20px_50px_rgba(16,185,129,0.15)]"
                    >
                      <span className="pointer-events-none absolute inset-0 bg-gradient-to-r from-emerald-500/15 via-transparent to-transparent opacity-0 transition-opacity group-hover:opacity-100" />
                      <div className="relative flex items-center gap-3 px-4 py-3.5">
                        <span className="flex size-10 flex-shrink-0 items-center justify-center rounded-full border border-emerald-500/20 bg-emerald-500/10 text-lg">
                          {example.icon}
                        </span>
                        <span className="text-sm font-medium text-gray-200 group-hover:text-white transition-colors">
                          {example.text}
                        </span>
                      </div>
                    </button>
                  ))}
                </div>
              </div>
            ) : (
              messages.map((msg, i) => (
                <ChatMessage key={i} role={msg.role} content={msg.content} />
              ))
            )}
          </div>

          {/* Input */}
          <div className="flex items-center gap-3">
            <div className="flex flex-1 items-center bg-white/5 border border-white/10 rounded-2xl overflow-hidden shadow-[0_15px_45px_rgba(0,0,0,0.35)]">
              <input
                ref={inputRef}
                type="text"
                value={input}
                onChange={(e) => setInput(e.target.value)}
                onKeyDown={(e) =>
                  e.key === "Enter" && !isGenerating && sendMessage(input)
                }
                disabled={isGenerating}
                className="flex-grow bg-transparent px-5 py-3.5 text-base text-white placeholder:text-gray-500 focus:outline-none disabled:opacity-40"
                placeholder="Ask Maincoder to write some code..."
              />
              <button
                onClick={() => sendMessage(input)}
                disabled={isGenerating || !input.trim()}
                className="h-full px-5 py-3.5 bg-emerald-600 hover:bg-emerald-500 disabled:bg-emerald-600/30 disabled:cursor-not-allowed text-white font-semibold transition-all"
              >
                <Send size={20} />
              </button>
            </div>
          </div>
        </div>
      </div>
    </div>
  );
};

export default App;