"use client";

import { useState, useEffect, useCallback, useMemo } from "react";
import { Card } from "@/components/ui/card";
import { Button } from "@/components/ui/button";
import { Flag, Clock, Hash, BarChart, ArrowRight, Bot } from "lucide-react";
import inference from "@/lib/inference";
import ReasoningTrace, { Run, Step } from "./reasoning-trace";
import ForceDirectedGraph from "./force-directed-graph";

const mockRun: Run = {
  steps: [
    {
      type: "start",
      article: "Dogs",
      metadata: {
        message: "Starting Node",
      },
    },
    {
      type: "step",
      article: "Dogs",
      links: ["Dogs", "Cats", "Birds"],
      metadata: {
        conversation: [
          {
            role: "user",
            content: "I want to go to the moon",
          },
          {
            role: "assistant",
            content: "I want to go to the moon",
          },
        ],
      },
    },
  ],
};
const buildPrompt = (
  current: string,
  target: string,
  path_so_far: string[],
  links: string[]
) => {
  const formatted_links = links
    .map((link, index) => `${index + 1}. ${link}`)
    .join("\n");
  const path_so_far_str = path_so_far.join(" -> ");

  return `You are playing WikiRun, trying to navigate from one Wikipedia article to another using only links.

IMPORTANT: You MUST put your final answer in <answer>NUMBER</answer> tags, where NUMBER is the link number.
For example, if you want to choose link 3, output <answer>3</answer>.

Current article: ${current}
Target article: ${target}
You have ${links.length} link(s) to choose from:
${formatted_links}

Your path so far: ${path_so_far_str}

Think about which link is most likely to lead you toward the target article.
First, analyze each link briefly and how it connects to your goal, then select the most promising one.

Remember to format your final answer by explicitly writing out the xml number tags like this: <answer>NUMBER</answer>`;
};

const API_BASE = "http://localhost:8000";

interface GameComponentProps {
  player: "me" | "model";
  model?: string;
  maxHops: number;
  startPage: string;
  targetPage: string;
  onReset: () => void;
  maxTokens: number;
  maxLinks: number;
}

export default function GameComponent({
  player,
  model,
  maxHops,
  startPage,
  targetPage,
  onReset,
  maxTokens,
  maxLinks,
}: GameComponentProps) {
  const [currentPage, setCurrentPage] = useState<string>(startPage);
  const [currentPageLinks, setCurrentPageLinks] = useState<string[]>([]);
  const [linksLoading, setLinksLoading] = useState<boolean>(false);
  const [hops, setHops] = useState<number>(0);
  const [timeElapsed, setTimeElapsed] = useState<number>(0);
  const [visitedNodes, setVisitedNodes] = useState<string[]>([startPage]);
  const [gameStatus, setGameStatus] = useState<"playing" | "won" | "lost">(
    "playing"
  );
  const [reasoningTrace, setReasoningTrace] = useState<Run | null>({
    start_article: startPage,
    destination_article: targetPage,
    steps: [
      {
        type: "start",
        article: startPage,
        metadata: {
          message: "Starting Node",
        },
      },
    ],
  });
  const [isModelThinking, setIsModelThinking] = useState<boolean>(false);

  const runs = useMemo(() => {
    return reasoningTrace ? [reasoningTrace] : [];
  }, [reasoningTrace]);

  const fetchCurrentPageLinks = useCallback(async () => {
    setLinksLoading(true);
    const response = await fetch(
      `${API_BASE}/get_article_with_links/${currentPage}`
    );
    const data = await response.json();
    setCurrentPageLinks(data.links.slice(0, maxLinks));
    setLinksLoading(false);
  }, [currentPage, maxLinks]);

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

  useEffect(() => {
    if (gameStatus === "playing") {
      const timer = setInterval(() => {
        setTimeElapsed((prev) => prev + 1);
      }, 1000);

      return () => clearInterval(timer);
    }
  }, [gameStatus]);

  // Check win condition
  useEffect(() => {
    if (currentPage === targetPage) {
      setGameStatus("won");
    } else if (hops >= maxHops) {
      setGameStatus("lost");
    }
  }, [currentPage, targetPage, hops, maxHops]);

  const addStepToReasoningTrace = (step: Step) => {
    setReasoningTrace((prev) => {
      if (!prev) return { steps: [step], start_article: startPage, destination_article: targetPage };
      return { steps: [...prev.steps, step], start_article: startPage, destination_article: targetPage };
    });
  };

  const handleLinkClick = (link: string, userClicked: boolean = true) => {
    if (gameStatus !== "playing") return;

    setCurrentPage(link);
    setHops((prev) => prev + 1);
    setVisitedNodes((prev) => [...prev, link]);
    if (userClicked) {
      addStepToReasoningTrace({
        type: "step",
        article: link,
        links: currentPageLinks,
        metadata: {
          message: "User clicked link",
        },
      });
    }
  };

  const makeModelMove = async () => {
    setIsModelThinking(true);

    const prompt = buildPrompt(
      currentPage,
      targetPage,
      visitedNodes,
      currentPageLinks
    );

    const modelResponse = await inference({
      apiKey:
        window.localStorage.getItem("huggingface_access_token") || undefined,
      model: model,
      prompt,
      maxTokens: maxTokens,
    });
    console.log("Model response", modelResponse.content);

    const answer = modelResponse.content?.match(/<answer>(.*?)<\/answer>/)?.[1];
    if (!answer) {
      console.error("No answer found in model response");
      return;
    }

    // try parsing the answer as an integer
    const answerInt = parseInt(answer);
    if (isNaN(answerInt)) {
      console.error("Invalid answer found in model response");
      return;
    }

    if (answerInt < 1 || answerInt > currentPageLinks.length) {
      console.error(
        "Selected link out of bounds",
        answerInt,
        "from ",
        currentPageLinks.length,
        "links"
      );
      return;
    }

    const selectedLink = currentPageLinks[answerInt - 1];

    console.log(
      "Model picked selectedLink",
      selectedLink,
      "from ",
      currentPageLinks
    );

    addStepToReasoningTrace({
      type: "step",
      article: selectedLink,
      links: currentPageLinks,
      metadata: {
        message: "Model picked link",
        conversation: [
          {
            role: "user",
            content: prompt,
          },
          {
            role: "assistant",
            content: modelResponse.content || "",
          },
        ],
      },
    });

    handleLinkClick(selectedLink, false);
    setIsModelThinking(false);
  };

  const handleGiveUp = () => {
    setGameStatus("lost");
  };

  const formatTime = (seconds: number) => {
    const mins = Math.floor(seconds / 60);
    const secs = seconds % 60;
    return `${mins}:${secs < 10 ? "0" : ""}${secs}`;
  };

  return (
    <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
      <Card className="p-4 flex col-span-2">
        <h2 className="text-xl font-bold">Game Status</h2>
        <div className="grid grid-cols-4 gap-4 mb-4">
          <div className="bg-muted/30 p-3 rounded-md">
            <div className="flex items-center gap-2 text-sm font-medium text-muted-foreground mb-1">
              <ArrowRight className="h-4 w-4" /> Current
            </div>
            <div className="font-medium">{currentPage}</div>
          </div>
          <div className="bg-muted/30 p-3 rounded-md">
            <div className="flex items-center gap-2 text-sm font-medium text-muted-foreground mb-1">
              <Flag className="h-4 w-4" /> Target
            </div>
            <div className="font-medium">{targetPage}</div>
          </div>

          <div className="bg-muted/30 p-3 rounded-md">
            <div className="flex items-center gap-2 text-sm font-medium text-muted-foreground mb-1">
              <Hash className="h-4 w-4" /> Hops
            </div>
            <div className="font-medium">
              {hops} / {maxHops}
            </div>
          </div>

          <div className="bg-muted/30 p-3 rounded-md">
            <div className="flex items-center gap-2 text-sm font-medium text-muted-foreground mb-1">
              <Clock className="h-4 w-4" /> Time
            </div>
            <div className="font-medium">{formatTime(timeElapsed)}</div>
          </div>
        </div>

        {player === "model" && (
          <div className="mb-4 bg-blue-50 border border-blue-200 rounded-md p-3">
            <div className="flex items-center gap-2">
              <Bot className="h-4 w-4 text-blue-500" />
              <span className="font-medium text-blue-700">
                {model} {isModelThinking ? "is playing..." : "is playing"}
              </span>
            </div>
          </div>
        )}
      </Card>
      {/* Left pane - Current page and available links */}
      <Card className="p-4 flex flex-col">
        <h2 className="text-xl font-bold">Available Links</h2>
        <div className="flex justify-between items-center mb-4">
          {gameStatus !== "playing" && (
            <Button onClick={onReset} variant="outline">
              New Game
            </Button>
          )}
        </div>

        {/* Wikipedia iframe (mocked) */}
        {/* <div className="bg-muted/30 rounded-md flex-1 mb-4 overflow-hidden">
          <div className="bg-white p-4 border-b">
            <h2 className="text-xl font-bold">{currentPage}</h2>
            <p className="text-sm text-muted-foreground">
              https://en.wikipedia.org/wiki/{currentPage.replace(/\s+/g, "_")}
            </p>
          </div>
          <div className="p-4">
            <p className="text-sm">
              This is a mock Wikipedia page for {currentPage}. In the actual
              implementation, this would be an iframe showing the real Wikipedia
              page.
            </p>
          </div>
        </div> */}

        {/* Available links */}
        {gameStatus === "playing" && (
          <>
            <div className="grid grid-cols-2 sm:grid-cols-3 gap-2 mb-4 overflow-y-auto max-h-[200px]">
              {currentPageLinks
                .sort((a, b) => a.localeCompare(b))
                .map((link) => (
                  <Button
                    key={link}
                    variant="outline"
                    size="sm"
                    className="justify-start overflow-hidden text-ellipsis whitespace-nowrap"
                    onClick={() => handleLinkClick(link)}
                    disabled={player === "model" || isModelThinking}
                  >
                    {link}
                  </Button>
                ))}
            </div>

            {player === "model" && (
              <Button
                onClick={makeModelMove}
                disabled={isModelThinking || linksLoading}
              >
                Make Move
              </Button>
            )}
          </>
        )}

        {player === "model" && isModelThinking && gameStatus === "playing" && (
          <div className="flex items-center gap-2 text-sm animate-pulse mb-4">
            <Bot className="h-4 w-4" />
            <span>{model} is thinking...</span>
          </div>
        )}

        {gameStatus === "playing" && player === "me" && (
          <Button
            onClick={handleGiveUp}
            variant="destructive"
            className="mt-auto"
          >
            Give Up
          </Button>
        )}

        {gameStatus === "won" && (
          <div className="bg-green-100 text-green-800 p-4 rounded-md mt-auto">
            <h3 className="font-bold">
              {player === "model" ? `${model} won!` : "You won!"}
            </h3>
            <p>
              {player === "model" ? "It" : "You"} reached {targetPage} in {hops}{" "}
              hops.
            </p>
          </div>
        )}

        {gameStatus === "lost" && (
          <div className="bg-red-100 text-red-800 p-4 rounded-md mt-auto">
            <h3 className="font-bold">Game Over</h3>
            <p>
              {player === "model" ? `${model} didn't` : "You didn't"} reach{" "}
              {targetPage} within {maxHops} hops.
            </p>
          </div>
        )}
      </Card>

      <Card className="p-4 flex flex-col max-h-[500px] overflow-y-auto">
        <ReasoningTrace run={reasoningTrace} />
      </Card>

      {/* Right pane - Game stats and graph */}
      <Card className="p-4 flex flex-col">
        <div className="flex-1 bg-muted/30 rounded-md p-4">
          <div className="flex items-center gap-2 text-sm font-medium text-muted-foreground mb-2">
            <BarChart className="h-4 w-4" /> Path Visualization
          </div>

          <ForceDirectedGraph runs={runs} runId={null} />
        </div>
      </Card>
    </div>
  );
}

