import React, { useCallback, useEffect, useMemo, useState } from "react";
import GuidelinePage from "./pages/GuidelinePage";
import TrainingPage from "./pages/TrainingPage";
import EvaluationPage from "./pages/EvaluationPage";
import RoundCompletePage from "./pages/RoundCompletePage";
import rawPrompts from "./data/prompts.json";
import { normalizePromptList } from "./lib/promptSchema";
import "./styles/guideline.css";

const API_BASE_URL = (
  import.meta.env.VITE_API_BASE_URL ||
  (import.meta.env.DEV ? "http://localhost:8000" : "")
).replace(/\/$/, "");

const STORAGE_KEYS = {
  annotatorName: "t2av_annotator_name",
  trainingPassed: "t2av_training_passed"
};

function readSession(key, fallback) {
  try {
    const value = sessionStorage.getItem(key);
    return value === null ? fallback : JSON.parse(value);
  } catch {
    return fallback;
  }
}

function writeSession(key, value) {
  try {
    sessionStorage.setItem(key, JSON.stringify(value));
  } catch {
    // sessionStorage unavailable (e.g. private browsing) - state just won't survive a refresh.
  }
}

const normalizedPrompts = normalizePromptList(rawPrompts);

export default function App() {
  // Gating note: there is no router and no URL for 'training'/'evaluation' -
  // the only way to reach them is through the state transitions below, so
  // there is nothing to type into an address bar to skip ahead.
  //
  // Stage itself is intentionally NOT restored from sessionStorage: every
  // fresh load (including a same-tab refresh) should land back on the
  // guideline page rather than silently resuming mid-training or
  // mid-evaluation. annotatorName and trainingPassed are still remembered,
  // so a returning annotator gets recognized and can skip the worked
  // example again (see the collision handling in GuidelinePage) - only the
  // "which page am I on" state resets.
  const [stage, setStage] = useState("guideline");
  const [annotatorName, setAnnotatorName] = useState(() => readSession(STORAGE_KEYS.annotatorName, ""));
  const [trainingPassed, setTrainingPassed] = useState(() =>
    readSession(STORAGE_KEYS.trainingPassed, false)
  );

  const [videoManifest, setVideoManifest] = useState(null);
  const [videoManifestError, setVideoManifestError] = useState(null);

  const [assignmentData, setAssignmentData] = useState(null);
  const [assignmentError, setAssignmentError] = useState(null);

  useEffect(() => {
    let cancelled = false;
    fetch(`${API_BASE_URL}/videos`)
      .then((res) => {
        if (!res.ok) throw new Error(`Video catalog request failed (status ${res.status})`);
        return res.json();
      })
      .then((data) => {
        if (!cancelled) setVideoManifest(data);
      })
      .catch((err) => {
        if (!cancelled) setVideoManifestError(err.message || "Failed to load videos");
      });
    return () => {
      cancelled = true;
    };
  }, []);

  useEffect(() => writeSession(STORAGE_KEYS.annotatorName, annotatorName), [annotatorName]);
  useEffect(() => writeSession(STORAGE_KEYS.trainingPassed, trainingPassed), [trainingPassed]);

  // Defensive fallback: the UI never lets you reach 'evaluation' without
  // trainingPassed, but if sessionStorage were hand-edited, fall back to
  // training rather than trusting the stored stage.
  useEffect(() => {
    if (stage === "evaluation" && !trainingPassed) {
      setStage("training");
    }
  }, [stage, trainingPassed]);
  const effectiveStage = stage === "evaluation" && !trainingPassed ? "training" : stage;

  // Each annotator gets a balanced 20-video round, not the full catalog.
  // The backend decides everything about round state (in progress / just
  // completed / time to auto-advance to the next round) - this call is
  // always the same "give me my current status" request, made automatically
  // on every evaluation-stage visit and again once EvaluationPage reports
  // the last video was saved, so the completed/code state is picked up
  // without any manual "start new round" action.
  const fetchAssignment = useCallback(() => {
    if (!annotatorName.trim()) return;
    fetch(`${API_BASE_URL}/assignment`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ annotator: annotatorName })
    })
      .then((res) => {
        if (!res.ok) throw new Error(`Assignment request failed (status ${res.status})`);
        return res.json();
      })
      .then((data) => setAssignmentData(data))
      .catch((err) => setAssignmentError(err.message || "Failed to load your assignment"));
  }, [annotatorName]);

  useEffect(() => {
    if (effectiveStage !== "evaluation") return;
    fetchAssignment();
  }, [effectiveStage, fetchAssignment]);

  // All resolved (itemId, model) videos, keyed by taskId - the full 200,
  // before per-annotator filtering.
  const videoEntriesByTaskId = useMemo(() => {
    if (!videoManifest) return null;
    return new Map(videoManifest.items.map((entry) => [`${entry.itemId}__${entry.model}`, entry]));
  }, [videoManifest]);

  const itemsById = useMemo(
    () => new Map(normalizedPrompts.map((item) => [item.itemId, item])),
    []
  );

  // The annotator's actual queue: their current round's video_ids (already
  // excludes anything they've completed before, in any round), resolved
  // against the live video catalog. An assigned id that no longer resolves
  // (e.g. a video went missing after assignment) is skipped gracefully,
  // matching the existing missing-video handling elsewhere.
  const tasks = useMemo(() => {
    if (!videoEntriesByTaskId || !assignmentData) return [];

    const built = [];
    assignmentData.video_ids.forEach((taskId) => {
      const entry = videoEntriesByTaskId.get(taskId);
      if (!entry) return;
      const item = itemsById.get(entry.itemId);
      if (!item) return;
      built.push({
        taskId,
        itemId: entry.itemId,
        model: entry.model,
        videoUrl: entry.videoUrl,
        variant: entry.variant,
        item
      });
    });
    return built;
  }, [videoEntriesByTaskId, itemsById, assignmentData]);

  const handleContinueToTraining = () => {
    if (!annotatorName.trim()) return;
    setStage("training");
  };

  // Offered on GuidelinePage only once handleGuidelineLookupResult has
  // already confirmed (via /assignment/lookup) that this name passed
  // training before - lets a returning annotator skip straight back into
  // their in-progress round instead of re-doing the worked example.
  const handleSkipTraining = () => {
    if (!annotatorName.trim() || !trainingPassed) return;
    setStage("evaluation");
  };

  // A round can only ever exist for a name that has already passed training
  // (POST /assignment is only ever called after trainingPassed is set, see
  // the evaluation-stage effect above) - so GET /assignment/lookup finding
  // an existing round is a reliable, read-only signal that this annotator
  // already passed training in an earlier session. GuidelinePage already
  // makes this exact call for its name-collision warning; this just listens
  // in on that same result instead of adding a second request.
  const handleGuidelineLookupResult = useCallback((data) => {
    if (data?.exists) setTrainingPassed(true);
  }, []);

  const handlePassTraining = () => {
    setTrainingPassed(true);
    setStage("evaluation");
  };

  if (effectiveStage === "training") {
    return <TrainingPage onPass={handlePassTraining} trainingPassed={trainingPassed} />;
  }

  if (effectiveStage === "evaluation") {
    if (videoManifestError) {
      return (
        <div className="t2av-shell">
          <div className="g-main">
            <div className="callout missing-videos-banner">
              <span className="k">Could not load videos</span>
              {videoManifestError}. Check that the backend is running and reachable at{" "}
              {API_BASE_URL || "(same origin)"}.
            </div>
          </div>
        </div>
      );
    }
    if (!videoManifest) {
      return (
        <div className="t2av-shell">
          <div className="g-main">
            <p>Loading videos…</p>
          </div>
        </div>
      );
    }
    if (assignmentError) {
      return (
        <div className="t2av-shell">
          <div className="g-main">
            <div className="callout missing-videos-banner">
              <span className="k">Could not load your assignment</span>
              {assignmentError}. Check that the backend is running and reachable at{" "}
              {API_BASE_URL || "(same origin)"}.
            </div>
          </div>
        </div>
      );
    }
    if (!assignmentData) {
      return (
        <div className="t2av-shell">
          <div className="g-main">
            <p>Loading your assignment…</p>
          </div>
        </div>
      );
    }
    if (assignmentData.status === "completed") {
      return (
        <RoundCompletePage
          annotatorName={annotatorName}
          roundNumber={assignmentData.round_number}
          completionCode={assignmentData.completion_code}
          videoCount={assignmentData.actual_size}
        />
      );
    }
    if (tasks.length === 0) {
      return (
        <div className="t2av-shell">
          <div className="g-main">
            <div className="callout missing-videos-banner">
              <span className="k">No videos available</span>
              The video catalog loaded, but none of your {assignmentData.video_ids.length} assigned
              videos matched a video in {normalizedPrompts.length} prompt items. Check
              HF_VIDEO_DATASET_REPO and the dataset's folder structure.
            </div>
          </div>
        </div>
      );
    }
    return (
      <EvaluationPage
        tasks={tasks}
        missingVideos={videoManifest.missing}
        annotatorName={annotatorName}
        apiBaseUrl={API_BASE_URL}
        onRoundFinished={fetchAssignment}
        savedAnnotations={assignmentData.saved_annotations || {}}
      />
    );
  }

  return (
    <GuidelinePage
      annotatorName={annotatorName}
      onChangeAnnotatorName={setAnnotatorName}
      onContinue={handleContinueToTraining}
      onSkipTraining={handleSkipTraining}
      apiBaseUrl={API_BASE_URL}
      onLookupResult={handleGuidelineLookupResult}
    />
  );
}