/**
 * FADA Mobile - Main application logic.
 * Manages UI state, chat interaction, and autonomous pipeline.
 */
import { loadModel, runInference, getLoadingStatus, getActiveDevice, testTextOnly } from "./model.js";
import { parseModelOutput, rescalePredictions } from "./output-parser.js";
import { parseIntent } from "./intent-parser.js";
import { parseInterpretationJson, extractClsLabel, mapInterpretationToClasses } from "./class-mapper.js";
import { drawAnnotations } from "./visualizer.js";
import {
  INTERPRET_PROMPT, CLASSIFY_PROMPT, DEFAULT_DETECT_CLASSES,
  MAX_NEW_TOKENS, TEMPERATURE,
} from "./constants.js";

// State
let currentImage = null;
let imageFile = null;
let interpretation = null;
let classification = null;
let detClasses = null;
let segClasses = null;
let isProcessing = false;

// DOM refs
const chatLog = document.getElementById("chat-log");
const userInput = document.getElementById("user-input");
const sendBtn = document.getElementById("send-btn");
const imageInput = document.getElementById("image-input");
const imagePreview = document.getElementById("image-preview");
const autoBtn = document.getElementById("auto-btn");
const loadBtn = document.getElementById("load-btn");
const deviceSelect = document.getElementById("device-select");
const statusEl = document.getElementById("status");
const progressBar = document.getElementById("progress-bar");
const progressText = document.getElementById("progress-text");

// ---------------------------------------------------------------------------
// Init
// ---------------------------------------------------------------------------

sendBtn.addEventListener("click", handleSend);
userInput.addEventListener("keydown", (e) => {
  if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); handleSend(); }
});
imageInput.addEventListener("change", handleImageUpload);
autoBtn.addEventListener("click", handleAutonomous);
loadBtn.addEventListener("click", handleLoadModel);

checkWebGPU();

async function checkWebGPU() {
  if (!navigator.gpu) {
    addMessage("system", "WebGPU NOT available. Select CPU (WASM) to proceed, or use Chrome/Edge 113+.");
    deviceSelect.value = "wasm";
    updateGpuInfo("No WebGPU");
    return;
  }
  try {
    // Request high-performance adapter to prefer NVIDIA discrete GPU over Intel integrated
    const adapter = await navigator.gpu.requestAdapter({
      powerPreference: "high-performance",
    });
    if (adapter) {
      const info = await adapter.requestAdapterInfo?.() || {};
      const desc = info.description || info.device || "unknown GPU";
      const vendor = info.vendor || "";
      const fullDesc = `${vendor} ${desc}`.trim();
      addMessage("system", `WebGPU available (high-performance): ${fullDesc}`);
      updateGpuInfo(fullDesc);
      console.log("[FADA] WebGPU adapter info:", JSON.stringify(info));
    } else {
      addMessage("system", "WebGPU adapter not found. Defaulting to CPU.");
      deviceSelect.value = "wasm";
      updateGpuInfo("No adapter");
    }
  } catch (e) {
    console.warn("[FADA] WebGPU check error:", e);
  }
}

function updateGpuInfo(text) {
  const el = document.getElementById("gpu-info");
  if (el) el.textContent = `GPU: ${text}`;
}

// ---------------------------------------------------------------------------
// Model loading
// ---------------------------------------------------------------------------

async function handleLoadModel() {
  if (getLoadingStatus() === "ready") {
    setStatus("Model already loaded");
    return;
  }
  const device = deviceSelect.value;
  loadBtn.disabled = true;
  deviceSelect.disabled = true;
  setStatus(`Loading model on ${device}...`);
  showProgress(true);

  let shaderTimer = null;
  let shaderStart = 0;

  try {
    shaderStart = Date.now();
    await loadModel((info) => {
      if (info.status === "downloading") {
        const pct = info.progress ? info.progress.toFixed(1) : "0";
        const file = info.file ? info.file.split("/").pop() : "";
        setProgress(info.progress || 0, `Downloading ${file} (${pct}%)`);
      } else if (info.status === "compiling" || info.note?.includes("compiling")) {
        setStatus(info.note || "Compiling GPU shaders...");
        setProgress(100, "Download complete - compiling shaders...");
        // Start animated compilation timer
        if (!shaderTimer) {
          shaderTimer = setInterval(() => {
            const elapsed = Math.floor((Date.now() - shaderStart) / 1000);
            const dots = '.'.repeat((elapsed % 3) + 1);
            setStatus(`Compiling GPU shaders${dots} (${elapsed}s elapsed, typically 2-5 min)`);
          }, 1000);
        }
      } else if (info.status === "loading") {
        setProgress(0, `Loading ${info.file}...`);
      } else if (info.status === "ready") {
        setProgress(100, "Model ready!");
      }
    }, device);

    if (shaderTimer) clearInterval(shaderTimer);
    setStatus(`Model loaded on ${device} - ready for inference`);
    showProgress(false);
    enableChat(true);
    document.getElementById('test-text-btn').disabled = false;
    if (device === "webgpu") {
      addMessage("system", "Model loaded on GPU. First inference may take 1-2 min for shader compilation.");
    } else {
      addMessage("system", "Model loaded on CPU. Inference will be slower but reliable.");
    }
  } catch (err) {
    if (shaderTimer) clearInterval(shaderTimer);
    setStatus(`Error loading model: ${err.message}`);
    showProgress(false);
    loadBtn.disabled = false;
    deviceSelect.disabled = false;
    addMessage("system", `Failed to load model: ${err.message}`);
  }
}

function enableChat(enabled) {
  sendBtn.disabled = !enabled;
  userInput.disabled = !enabled;
  autoBtn.disabled = !enabled;
  if (enabled) userInput.focus();
}

// ---------------------------------------------------------------------------
// Image handling
// ---------------------------------------------------------------------------

function handleImageUpload(e) {
  const file = e.target.files[0];
  if (!file) return;
  imageFile = file;
  const img = new Image();
  img.onload = () => {
    const maxSide = 1024;
    let w = img.width, h = img.height;
    if (Math.max(w, h) > maxSide) {
      const scale = maxSide / Math.max(w, h);
      w = Math.round(w * scale);
      h = Math.round(h * scale);
    }
    const canvas = document.createElement("canvas");
    canvas.width = w; canvas.height = h;
    canvas.getContext("2d").drawImage(img, 0, 0, w, h);
    const resized = new Image();
    resized.onload = () => {
      currentImage = resized;
      imagePreview.innerHTML = "";
      imagePreview.appendChild(resized.cloneNode());
      interpretation = null;
      classification = null;
      detClasses = null;
      segClasses = null;
      addMessage("system", `Image uploaded (${w}x${h}). Type "interpret", "detect", or "segment" in the chat.`);
    };
    resized.src = canvas.toDataURL();
  };
  img.src = URL.createObjectURL(file);
}

// ---------------------------------------------------------------------------
// Chat handling
// ---------------------------------------------------------------------------

async function handleSend() {
  const text = userInput.value.trim();
  if (!text && !currentImage) return;
  if (isProcessing) return;
  if (getLoadingStatus() !== "ready") {
    addMessage("system", "Please load the model first (click 'Load Model').");
    return;
  }

  addMessage("user", text || "(image uploaded)");
  userInput.value = "";
  isProcessing = true;
  sendBtn.disabled = true;

  try {
    await processChatTurn(text);
  } catch (err) {
    console.error("Inference error:", err);
    removeThinking();
    addMessage("system", `Error: ${err.message}`);
  }

  isProcessing = false;
  sendBtn.disabled = false;
  userInput.focus();
}

async function processChatTurn(userText) {
  const intent = parseIntent(userText, !!currentImage, !!interpretation);

  if (intent.task === "greeting") {
    addMessage("assistant", "Welcome to FADA Mobile! Upload a fetal ultrasound image and I'll provide clinical interpretation. You can then ask me to detect or segment specific anatomical structures.");
    return;
  }
  if (intent.task === "general") {
    addMessage("assistant", "I can help with fetal ultrasound analysis. Try:\n- Upload an image and ask me to **interpret** it\n- Ask me to **detect** specific structures (e.g., 'detect brain')\n- Ask me to **segment** anatomy (e.g., 'segment cardiac')");
    return;
  }
  if (!currentImage) {
    addMessage("assistant", "Please upload a fetal ultrasound image first.");
    return;
  }

  if (intent.task === "interpret") {
    addMessage("assistant", "Analyzing image...", "thinking");
    setStatus("Running interpretation...");
    const t0 = performance.now();
    let tokensSoFar = 0;
    const rawInterp = await runInference(currentImage, INTERPRET_PROMPT, {
      temperature: TEMPERATURE,
      onToken: (tok, n) => {
        tokensSoFar = n;
        setStatus(`Interpreting... (${n} tokens generated)`);
      },
    });
    const interpTime = ((performance.now() - t0) / 1000).toFixed(1);
    interpretation = parseInterpretationJson(rawInterp);

    setStatus("Running classification...");
    const rawCls = await runInference(currentImage, CLASSIFY_PROMPT, {
      maxNewTokens: 256,
      temperature: TEMPERATURE,
    });
    const clsResult = parseInterpretationJson(rawCls);
    classification = extractClsLabel(clsResult);

    if (interpretation._parseSuccess) {
      const mapped = mapInterpretationToClasses(interpretation, classification);
      detClasses = mapped.det;
      segClasses = mapped.seg;
    }

    removeThinking();
    let response = formatInterpretation(interpretation);
    if (classification) response += `\n**Classification:** ${classification}`;
    response += `\n\n*${tokensSoFar} tokens in ${interpTime}s on ${getActiveDevice()}*`;
    response += "\n\n---\nYou can now ask me to **detect** or **segment** specific structures.";
    addMessage("assistant", response);
    setStatus("Ready");
    return;
  }

  if (intent.task === "classify") {
    addMessage("assistant", "Classifying...", "thinking");
    setStatus("Running classification...");
    const rawCls = await runInference(currentImage, CLASSIFY_PROMPT, { maxNewTokens: 256, temperature: TEMPERATURE });
    const clsResult = parseInterpretationJson(rawCls);
    classification = extractClsLabel(clsResult);
    removeThinking();
    addMessage("assistant", `**Classification:** ${classification || "Unknown"}`);
    setStatus("Ready");
    return;
  }

  if (intent.task === "detect") {
    const classes = intent.classes || detClasses || DEFAULT_DETECT_CLASSES;
    const prompt = buildDetectPrompt(classes);
    addMessage("assistant", `Detecting: ${classes}...`, "thinking");
    setStatus("Running detection...");
    const raw = await runInference(currentImage, prompt, {
      temperature: TEMPERATURE,
      onToken: (tok, n) => setStatus(`Detecting... (${n} tokens)`),
    });
    removeThinking();
    showAnnotation(raw, "Detection");
    setStatus("Ready");
    return;
  }

  if (intent.task === "segment") {
    const classes = intent.segClasses || intent.classes || segClasses || detClasses || DEFAULT_DETECT_CLASSES;
    const prompt = buildSegmentPrompt(classes);
    addMessage("assistant", `Segmenting: ${classes}...`, "thinking");
    setStatus("Running segmentation...");
    const raw = await runInference(currentImage, prompt, {
      temperature: TEMPERATURE,
      onToken: (tok, n) => setStatus(`Segmenting... (${n} tokens)`),
    });
    removeThinking();
    showAnnotation(raw, "Segmentation");
    setStatus("Ready");
    return;
  }
}

// ---------------------------------------------------------------------------
// Autonomous pipeline
// ---------------------------------------------------------------------------

async function handleAutonomous() {
  if (!currentImage) {
    addMessage("system", "Please upload an image first.");
    return;
  }
  if (getLoadingStatus() !== "ready") {
    addMessage("system", "Please load the model first.");
    return;
  }
  if (isProcessing) return;
  isProcessing = true;
  autoBtn.disabled = true;
  sendBtn.disabled = true;

  try {
    addMessage("system", "Starting autonomous 5-phase pipeline...");

    // Phase 1: Interpret
    setStatus("Phase 1/5: Interpreting...");
    addMessage("assistant", "Phase 1: Interpreting image...", "thinking");
    const t0 = performance.now();
    const rawInterp = await runInference(currentImage, INTERPRET_PROMPT, {
      temperature: TEMPERATURE,
      onToken: (tok, n) => setStatus(`Phase 1/5: Interpreting... (${n} tokens)`),
    });
    const interpTime = ((performance.now() - t0) / 1000).toFixed(1);
    interpretation = parseInterpretationJson(rawInterp);
    removeThinking();
    addMessage("assistant", formatInterpretation(interpretation) + `\n*${interpTime}s*`);

    // Phase 2: Classify
    setStatus("Phase 2/5: Classifying...");
    const rawCls = await runInference(currentImage, CLASSIFY_PROMPT, { maxNewTokens: 256, temperature: TEMPERATURE });
    const clsResult = parseInterpretationJson(rawCls);
    classification = extractClsLabel(clsResult);
    addMessage("assistant", `**Classification:** ${classification || "Unknown"}`);

    // Phase 3: Map
    setStatus("Phase 3/5: Mapping classes...");
    let mapped;
    if (interpretation._parseSuccess) {
      mapped = mapInterpretationToClasses(interpretation, classification);
    } else {
      mapped = { det: DEFAULT_DETECT_CLASSES, seg: null, tier: "fallback" };
    }
    detClasses = mapped.det;
    segClasses = mapped.seg;
    addMessage("system", `Mapped: detect=[${mapped.det}], segment=[${mapped.seg || "none"}], tier=${mapped.tier}`);

    // Phase 4: Detect
    setStatus("Phase 4/5: Detecting...");
    addMessage("assistant", `Detecting: ${detClasses}...`, "thinking");
    const rawDet = await runInference(currentImage, buildDetectPrompt(detClasses), {
      temperature: TEMPERATURE,
      onToken: (tok, n) => setStatus(`Phase 4/5: Detecting... (${n} tokens)`),
    });
    removeThinking();
    showAnnotation(rawDet, "Detection");

    // Phase 5: Segment (conditional)
    if (segClasses) {
      setStatus("Phase 5/5: Segmenting...");
      addMessage("assistant", `Segmenting: ${segClasses}...`, "thinking");
      const rawSeg = await runInference(currentImage, buildSegmentPrompt(segClasses), {
        temperature: TEMPERATURE,
        onToken: (tok, n) => setStatus(`Phase 5/5: Segmenting... (${n} tokens)`),
      });
      removeThinking();
      showAnnotation(rawSeg, "Segmentation");
    } else {
      addMessage("system", "Phase 5: Segmentation skipped (no classes mapped).");
    }

    setStatus("Pipeline complete!");
    addMessage("system", "Autonomous pipeline complete. You can ask follow-up questions.");
  } catch (err) {
    removeThinking();
    console.error("Pipeline error:", err);
    addMessage("system", `Pipeline error: ${err.message}`);
    setStatus("Error");
  }

  isProcessing = false;
  autoBtn.disabled = false;
  sendBtn.disabled = false;
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function buildDetectPrompt(classes) {
  return `Detect all instances of the following anatomical structures in this fetal ultrasound image: ${classes}. For each detection, return a JSON object with "bbox_2d" as [x_min, y_min, x_max, y_max] in normalized 0-1000 coordinates and "label" as the structure name.`;
}

function buildSegmentPrompt(classes) {
  return `Detect and segment all instances of the following anatomical structures in this fetal ultrasound image: ${classes}. For each instance, return a JSON object with "bbox_2d" as [x_min, y_min, x_max, y_max] in normalized 0-1000 coordinates, "label" as the structure name, and "mask" as a list of [x, y] polygon vertices in normalized 0-1000 coordinates.`;
}

function formatInterpretation(interp) {
  if (!interp._parseSuccess) {
    return `**Interpretation (raw):**\n\`\`\`\n${(interp._rawText || "").slice(0, 1000)}\n\`\`\``;
  }
  const fields = [
    ["Anatomical Structures", "anatomical_structures"],
    ["Fetal Orientation", "fetal_orientation"],
    ["Imaging Plane", "imaging_plane"],
    ["Biometric Measurements", "biometric_measurements"],
    ["Gestational Age", "gestational_age"],
    ["Image Quality", "image_quality"],
    ["Normality Assessment", "normality_assessment"],
    ["Clinical Recommendations", "clinical_recommendations"],
  ];
  let parts = "**Clinical Interpretation:**\n\n";
  for (const [title, key] of fields) {
    let value = interp[key] || "N/A";
    if (typeof value === "object") value = JSON.stringify(value, null, 2);
    parts += `**${title}:** ${value}\n\n`;
  }
  return parts;
}

function showAnnotation(rawText, taskName) {
  try {
    const parsed = parseModelOutput(rawText);
    const rescaled = rescalePredictions(parsed, currentImage.width, currentImage.height);
    const nDet = rescaled.detections.length;
    const nSeg = rescaled.segmentations.length;
    const nKp = rescaled.keypoints.length;

    if (nDet + nSeg + nKp === 0) {
      addMessage("assistant", `**${taskName}:** No structures found for the requested classes.`);
      return;
    }

    const canvas = document.createElement("canvas");
    drawAnnotations(canvas, currentImage, rescaled);
    const imgEl = new Image();
    imgEl.src = canvas.toDataURL();
    imgEl.className = "annotation-img";

    const labels = new Set();
    rescaled.detections.forEach(d => labels.add(d.label));
    rescaled.segmentations.forEach(s => labels.add(s.label));
    rescaled.keypoints.forEach(k => labels.add(k.label));

    const parts = [];
    if (nDet > 0) parts.push(`${nDet} detections`);
    if (nSeg > 0) parts.push(`${nSeg} masks`);
    if (nKp > 0) parts.push(`${nKp} keypoints`);

    addMessage("assistant", `**${taskName} results** (${parts.join(", ")}: ${[...labels].sort().join(", ")})`);
    chatLog.lastElementChild.appendChild(imgEl);
  } catch (err) {
    addMessage("assistant", `**${taskName}** completed but annotation failed. Raw output:\n\`\`\`\n${rawText.slice(0, 500)}\n\`\`\``);
  }
}

function addMessage(role, text, extraClass = "") {
  const div = document.createElement("div");
  div.className = `message ${role} ${extraClass}`.trim();
  let html = text
    .replace(/\*\*(.*?)\*\*/g, "<strong>$1</strong>")
    .replace(/\n/g, "<br>");
  div.innerHTML = html;
  chatLog.appendChild(div);
  chatLog.scrollTop = chatLog.scrollHeight;
}

function removeThinking() {
  const el = chatLog.querySelector(".thinking");
  if (el) el.remove();
}

function setStatus(text) {
  statusEl.textContent = text;
}

function showProgress(show) {
  document.getElementById("progress-container").style.display = show ? "block" : "none";
}

function setProgress(pct, text) {
  progressBar.style.width = `${pct}%`;
  progressText.textContent = text || "";
}

// ---------------------------------------------------------------------------
// Text-only diagnostic test
// ---------------------------------------------------------------------------

window.runTextTest = async () => {
  const btn = document.getElementById('test-text-btn');
  btn.disabled = true;
  btn.textContent = "Testing...";

  try {
    const result = await testTextOnly();
    // Display result in chat
    addMessage("assistant", `<strong>Diagnostic:</strong> ${result}`);
    alert(result);
  } catch(e) {
    alert("Test failed: " + e.message);
    console.error("[FADA-TEST] Error:", e);
    addMessage("system", `Text-only test error: ${e.message}`);
  } finally {
    btn.disabled = false;
    btn.textContent = "Test Text-Only";
  }
};