server.ts

import express from "express";
import cors from "cors";
import path from "path";
import dotenv from "dotenv";
import { GoogleGenAI, Type } from "@google/genai";
import { createServer as createViteServer } from "vite";

dotenv.config();

const PORT = 3000;

interface RoadmapWeek {
  number: number;
  title: string;
  focus: string;
  steps: string[];
  tools: string[];
}

interface RoadmapResponse {
  subtitle: string;
  weeks: RoadmapWeek[];
}

function getWeekCount(dur: string): number {
  const lower = (dur || "").toLowerCase();
  if (lower.includes("4")) return 4;
  if (lower.includes("8")) return 8;
  if (lower.includes("12")) return 12;
  if (lower.includes("6") || lower.includes("month")) return 16;
  const parsed = parseInt(dur, 10);
  return isNaN(parsed) ? 4 : parsed;
}

function cleanJsonOutput(raw: string): string {
  let cleaned = raw.trim();
  if (cleaned.startsWith("```json")) {
    cleaned = cleaned.replace(/^```json\s*/, "").replace(/\s*```$/, "");
  } else if (cleaned.startsWith("```")) {
    cleaned = cleaned.replace(/^```\s*/, "").replace(/\s*```$/, "");
  }
  return cleaned;
}

// Default Groq model managed entirely server-side
const GROQ_FALLBACK_MODELS = [
  "llama-3.3-70b-versatile",
  "llama-3.1-8b-instant",
  "deepseek-r1-distill-llama-70b",
  "gemma2-9b-it",
];

async function getAvailableGroqModels(apiKey: string): Promise<string[]> {
  try {
    const res = await fetch("https://api.groq.com/openai/v1/models", {
      headers: { Authorization: `Bearer ${apiKey}` },
    });
    if (res.ok) {
      const data = await res.json();
      if (Array.isArray(data?.data)) {
        return data.data
          .map((m: any) => m.id)
          .filter(
            (id: string) =>
              id &&
              !id.includes("whisper") &&
              !id.includes("guard") &&
              !id.includes("vision") &&
              !id.includes("tts")
          );
      }
    }
  } catch (e) {
    console.warn("Could not query Groq models list:", e);
  }
  return [];
}

async function generateWithGroq(apiKey: string, prompt: string): Promise<string> {
  const discovered = await getAvailableGroqModels(apiKey);
  const candidateModels = Array.from(new Set([...discovered, ...GROQ_FALLBACK_MODELS]));

  let lastError: Error | null = null;

  for (const model of candidateModels) {
    try {
      const response = await fetch("https://api.groq.com/openai/v1/chat/completions", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${apiKey}`,
        },
        body: JSON.stringify({
          model,
          messages: [
            {
              role: "system",
              content:
                "You are a specialized JSON-only roadmap generation API. You strictly return parseable JSON matching the requested structure without code fences or explanations.",
            },
            {
              role: "user",
              content: prompt,
            },
          ],
          temperature: 0.7,
          response_format: { type: "json_object" },
        }),
      });

      if (!response.ok) {
        const errorText = await response.text();
        console.warn(`Groq model '${model}' returned HTTP ${response.status}: ${errorText}`);
        lastError = new Error(`Groq API returned HTTP ${response.status}: ${errorText}`);
        continue;
      }

      const data = await response.json();
      const content = data?.choices?.[0]?.message?.content || "";
      if (content) {
        return content;
      }
    } catch (err: any) {
      lastError = err;
      continue;
    }
  }

  throw lastError || new Error("Failed to generate roadmap using Groq API");
}

async function getAvailableGrokModels(apiKey: string): Promise<string[]> {
  try {
    const response = await fetch("https://api.x.ai/v1/models", {
      headers: {
        Authorization: `Bearer ${apiKey}`,
      },
    });
    if (response.ok) {
      const data = await response.json();
      if (Array.isArray(data?.data)) {
        return data.data.map((m: any) => m.id).filter(Boolean);
      }
    }
  } catch (err) {
    console.warn("Could not query xAI models endpoint:", err);
  }
  return [];
}

async function generateWithGrok(apiKey: string, prompt: string): Promise<string> {
  const customModel = (process.env.GROK_MODEL || "").trim();
  const discoveredModels = await getAvailableGrokModels(apiKey);
  
  const candidateModels = Array.from(
    new Set([
      ...(customModel ? [customModel] : []),
      ...discoveredModels,
      "grok-2-1212",
      "grok-2",
      "grok-beta",
      "grok-3",
      "grok-4",
      "grok-2-vision-1212",
    ])
  );

  let lastError: Error | null = null;

  for (const model of candidateModels) {
    try {
      const response = await fetch("https://api.x.ai/v1/chat/completions", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${apiKey}`,
        },
        body: JSON.stringify({
          model,
          messages: [
            {
              role: "system",
              content:
                "You are a specialized JSON-only roadmap generation API. You strictly return parseable JSON matching the requested structure without code fences or explanations.",
            },
            {
              role: "user",
              content: prompt,
            },
          ],
          temperature: 0.7,
        }),
      });

      if (!response.ok) {
        const errorText = await response.text();
        if (
          response.status === 400 &&
          (errorText.includes("Model not found") ||
            errorText.includes("model_not_found") ||
            errorText.includes("invalid-argument") ||
            errorText.includes("does not exist"))
        ) {
          lastError = new Error(`xAI Grok model ${model} unavailable: ${errorText}`);
          continue;
        }
        throw new Error(`xAI Grok API returned HTTP ${response.status}: ${errorText}`);
      }

      const data = await response.json();
      const content = data?.choices?.[0]?.message?.content || "";
      if (content) {
        return content;
      }
    } catch (err: any) {
      lastError = err;
      if (err.message && (err.message.includes("Model not found") || err.message.includes("unavailable"))) {
        continue;
      }
      throw err;
    }
  }

  throw lastError || new Error("Failed to generate roadmap using xAI Grok API");
}

async function generateWithOpenRouter(apiKey: string, prompt: string): Promise<string> {
  const customModel = (process.env.OPENROUTER_MODEL || "").trim();
  const candidateModels = customModel
    ? [customModel, "x-ai/grok-2-1212", "x-ai/grok-beta", "google/gemini-2.5-flash"]
    : ["x-ai/grok-2-1212", "x-ai/grok-beta", "google/gemini-2.5-flash"];

  let lastError: Error | null = null;

  for (const model of candidateModels) {
    try {
      const response = await fetch("https://openrouter.ai/api/v1/chat/completions", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Authorization: `Bearer ${apiKey}`,
          "HTTP-Referer": "http://localhost:3000",
          "X-Title": "PathAI",
        },
        body: JSON.stringify({
          model,
          messages: [
            {
              role: "system",
              content:
                "You are a specialized JSON-only roadmap generation API. You strictly return parseable JSON matching the requested structure without code fences or explanations.",
            },
            {
              role: "user",
              content: prompt,
            },
          ],
          temperature: 0.7,
          response_format: { type: "json_object" },
        }),
      });

      if (!response.ok) {
        const errorText = await response.text();
        if (response.status === 400 || response.status === 404) {
          lastError = new Error(`OpenRouter model ${model} unavailable: ${errorText}`);
          continue;
        }
        throw new Error(`OpenRouter API returned HTTP ${response.status}: ${errorText}`);
      }

      const data = await response.json();
      const content = data?.choices?.[0]?.message?.content || "";
      if (content) {
        return content;
      }
    } catch (err: any) {
      lastError = err;
      continue;
    }
  }

  throw lastError || new Error("Failed to generate roadmap using OpenRouter API");
}

async function generateWithGemini(apiKey: string, prompt: string): Promise<string> {
  const ai = new GoogleGenAI({
    apiKey,
    httpOptions: {
      headers: {
        "User-Agent": "aistudio-build",
      },
    },
  });

  const response = await ai.models.generateContent({
    model: "gemini-3.6-flash",
    contents: prompt,
    config: {
      systemInstruction:
        "You are a specialized JSON-only roadmap generation API. You strictly return parseable JSON matching the requested structure without code fences or markdown.",
      responseMimeType: "application/json",
      responseSchema: {
        type: Type.OBJECT,
        properties: {
          subtitle: {
            type: Type.STRING,
            description: "A brief poetic sentence about this specific journey (under 20 words)",
          },
          weeks: {
            type: Type.ARRAY,
            description: "Sequential list of weekly milestone cards",
            items: {
              type: Type.OBJECT,
              properties: {
                number: {
                  type: Type.INTEGER,
                  description: "Week number starting at 1",
                },
                title: {
                  type: Type.STRING,
                  description: "Short, evocative title for this week (3-6 words)",
                },
                focus: {
                  type: Type.STRING,
                  description: "One sentence outlining focus and practical progress",
                },
                steps: {
                  type: Type.ARRAY,
                  items: { type: Type.STRING },
                  description: "3 concrete action items for this week",
                },
                tools: {
                  type: Type.ARRAY,
                  items: { type: Type.STRING },
                  description: "2 recommended tools or resources",
                },
              },
              required: ["number", "title", "focus", "steps", "tools"],
            },
          },
        },
        required: ["subtitle", "weeks"],
      },
    },
  });

  const text = response.text;
  if (!text) {
    throw new Error("Empty response received from Gemini API");
  }
  return text;
}

function resolveProvider(): { provider: "groq" | "grok" | "openrouter" | "gemini"; apiKey: string } {
  const reqProvider = (process.env.API_PROVIDER || "").trim().toLowerCase();
  const groqKey = (process.env.GROQ_API_KEY || "").trim();
  const grokKey = (process.env.GROK_API_KEY || process.env.XAI_API_KEY || "").trim();
  const openrouterKey = (process.env.OPENROUTER_API_KEY || "").trim();
  const geminiKey = (process.env.GEMINI_API_KEY || "").trim();

  // If a key starts with 'gsk_', it's definitely a Groq (console.groq.com) key
  if (groqKey.startsWith("gsk_")) {
    return { provider: "groq", apiKey: groqKey };
  }
  if (grokKey.startsWith("gsk_")) {
    return { provider: "groq", apiKey: grokKey };
  }

  // If a key starts with 'xai-', it's definitely an xAI Grok (console.x.ai) key
  if (grokKey.startsWith("xai-") || grokKey) {
    if (reqProvider === "grok" || reqProvider === "xai" || (!groqKey && !geminiKey && !openrouterKey)) {
      return { provider: "grok", apiKey: grokKey };
    }
  }

  // 1. Honor explicit API_PROVIDER choice
  if ((reqProvider === "groq" || reqProvider === "groq.com") && (groqKey || grokKey)) {
    return { provider: "groq", apiKey: groqKey || grokKey };
  }
  if ((reqProvider === "grok" || reqProvider === "xai" || reqProvider === "x.ai") && grokKey) {
    return { provider: "grok", apiKey: grokKey };
  }
  if (reqProvider === "openrouter" && openrouterKey) {
    return { provider: "openrouter", apiKey: openrouterKey };
  }
  if (reqProvider === "gemini" && geminiKey) {
    return { provider: "gemini", apiKey: geminiKey };
  }

  // 2. If specified without key
  if (reqProvider === "groq") {
    return { provider: "groq", apiKey: groqKey || grokKey };
  }
  if (reqProvider === "grok" || reqProvider === "xai") {
    return { provider: "grok", apiKey: grokKey };
  }
  if (reqProvider === "openrouter") {
    return { provider: "openrouter", apiKey: openrouterKey };
  }
  if (reqProvider === "gemini") {
    return { provider: "gemini", apiKey: geminiKey };
  }

  // 3. Automatic selection based on available keys
  if (groqKey) return { provider: "groq", apiKey: groqKey };
  if (grokKey) return { provider: "grok", apiKey: grokKey };
  if (openrouterKey) return { provider: "openrouter", apiKey: openrouterKey };
  return { provider: "gemini", apiKey: geminiKey };
}

async function handleRoadmapRequest(req: express.Request, res: express.Response) {
  try {
    const { from, to, dur } = req.body || {};

    if (!from || !to || !dur) {
      res.status(400).json({ error: "Missing required fields: from, to, dur" });
      return;
    }

    const { provider, apiKey } = resolveProvider();

    if (!apiKey) {
      console.error(`API key is missing for provider '${provider}'`);
      res.status(400).json({
        error: `No API key found for provider '${provider}'. Please add your API key in Settings (GROQ_API_KEY, GROK_API_KEY/XAI_API_KEY from console.x.ai, or GEMINI_API_KEY).`,
      });
      return;
    }

    const weekCount = getWeekCount(dur);
    const prompt = `You are a journey architect. Create a cinematic, personal roadmap for someone going from "${from}" to "${to}" over ${dur}.

Return ONLY valid JSON (no markdown code blocks, no backticks, no explanations, no preamble).
Strictly match this JSON schema:
{
  "subtitle": "A brief poetic sentence about this specific journey (under 20 words)",
  "weeks": [
    {
      "number": 1,
      "title": "Poetic evocative title for this phase",
      "focus": "One sentence describing the emotional or practical focus of this week",
      "steps": [
        "Concrete action 1",
        "Concrete action 2",
        "Concrete action 3"
      ],
      "tools": [
        "Tool 1",
        "Tool 2"
      ]
    }
  ]
}

Generate exactly ${weekCount} weeks numbered 1 to ${weekCount}.
Requirements:
- subtitle: string
- weeks: array of exactly ${weekCount} objects
- each week must have:
  - "number": integer (1 to ${weekCount})
  - "title": string (poetic and evocative)
  - "focus": string (emotional/practical focus)
  - "steps": array of exactly 3 actionable strings
  - "tools": array of exactly 2 useful tools
Do not return markdown, backticks, or any conversational text. Return only valid parseable JSON.`;

    let rawOutput = "";
    try {
      if (provider === "groq") {
        rawOutput = await generateWithGroq(apiKey, prompt);
      } else if (provider === "grok") {
        rawOutput = await generateWithGrok(apiKey, prompt);
      } else if (provider === "openrouter") {
        rawOutput = await generateWithOpenRouter(apiKey, prompt);
      } else {
        rawOutput = await generateWithGemini(apiKey, prompt);
      }
    } catch (apiErr: any) {
      console.error(`Provider '${provider}' generation error:`, apiErr);
      res.status(502).json({
        error: apiErr?.message || `Failed to generate roadmap from ${provider} API.`,
      });
      return;
    }

    const cleaned = cleanJsonOutput(rawOutput);
    let parsed: RoadmapResponse;
    try {
      parsed = JSON.parse(cleaned);
    } catch (parseErr: any) {
      console.error("JSON parsing failed on AI output:", rawOutput);
      res.status(502).json({
        error: "AI provider returned invalid JSON format. Please try again.",
      });
      return;
    }

    if (!parsed || !parsed.subtitle || !Array.isArray(parsed.weeks)) {
      res.status(502).json({
        error: "AI provider returned an incomplete roadmap structure. Please try again.",
      });
      return;
    }

    parsed.weeks = parsed.weeks.slice(0, weekCount).map((w, idx) => ({
      number: typeof w.number === "number" ? w.number : idx + 1,
      title: w.title || `Phase ${idx + 1}`,
      focus: w.focus || "Dedicate focus to your journey goals.",
      steps:
        Array.isArray(w.steps) && w.steps.length === 3
          ? w.steps
          : Array.isArray(w.steps)
          ? w.steps.slice(0, 3)
          : [
              "Research core principles",
              "Apply your learning",
              "Reflect on progress",
            ],
      tools:
        Array.isArray(w.tools) && w.tools.length === 2
          ? w.tools
          : Array.isArray(w.tools)
          ? w.tools.slice(0, 2)
          : ["Visual Studio Code", "Notion"],
    }));

    res.json(parsed);
  } catch (error: any) {
    console.error("Roadmap request handler error:", error);
    res.status(500).json({
      error: error.message || "Failed to generate roadmap.",
    });
  }
}

async function startServer() {
  const app = express();

  // Middleware
  app.use(cors());
  app.use(express.json());

  // Handle JSON body parse errors with JSON responses
  app.use((err: any, _req: express.Request, res: express.Response, next: express.NextFunction) => {
    if (err) {
      res.status(400).json({ error: "Invalid JSON payload in request." });
      return;
    }
    next();
  });

  // API Routes FIRST
  app.post("/roadmap", handleRoadmapRequest);
  app.post("/api/roadmap", handleRoadmapRequest);

  // Health check
  app.get("/api/health", (_req, res) => {
    const { provider } = resolveProvider();
    res.json({
      status: "ok",
      service: "PathAI Backend",
      provider,
      configuredProvider: process.env.API_PROVIDER || "auto",
    });
  });

  // Specific 404 handler for API routes
  app.all("/api/*", (_req, res) => {
    res.status(404).json({ error: "API endpoint not found" });
  });

  // Global error handler for API
  app.use((err: any, _req: express.Request, res: express.Response, _next: express.NextFunction) => {
    console.error("Unhandled server error:", err);
    res.status(err.status || 500).json({
      error: err.message || "Internal server error",
    });
  });

  // Vite middleware for development / static serving in production
  if (process.env.NODE_ENV !== "production") {
    const vite = await createViteServer({
      server: { middlewareMode: true },
      appType: "spa",
    });
    app.use(vite.middlewares);
  } else {
    const distPath = path.join(process.cwd(), "dist");
    app.use(express.static(distPath));
    app.get("*", (_req, res) => {
      res.sendFile(path.join(distPath, "index.html"));
    });
  }

  app.listen(PORT, "0.0.0.0", () => {
    console.log(`PathAI server running on http://0.0.0.0:${PORT}`);
  });
}

startServer();