tesla-roadtrip/server/services/llm/GrokHeadlessClient.ts

/**
 * Tesla Roadtrip — Grok Headless Client (with real xAI API fallback)
 * Maximum logging + strict structured output for map rendering
 */
import { spawn } from 'child_process';
import { mkdtemp, rm } from 'fs/promises';
import { tmpdir } from 'os';
import { join } from 'path';
import { createLogger } from '../../lib/logger.js';
import { env } from '../../config/env.js';
import crypto from 'crypto';

const log = createLogger('grok-headless');
const SENTINEL = 'ITINERARY_UPDATE:';

export interface ChatMessage { role: 'user' | 'assistant' | 'system'; content: string; }
export interface GrokResponse { text: string; updatedItinerary?: any; }

export class GrokHeadlessClient {
  private useFallback = !!env.xaiApiKey || !env.grokEnabled;

  private buildPrompt(messages: ChatMessage[], itinerary: any, vehicle: string) {
    return `You are Grok Drive — an expert Tesla road trip planner for the UK and Europe.

Current vehicle: ${vehicle}
Current itinerary: ${JSON.stringify(itinerary || {}, null, 2)}

CRITICAL OUTPUT RULES (must follow exactly):
- When the user gives a clear origin and destination, immediately create a reasonable first-draft multi-day itinerary.
- For every stop in the itinerary (Superchargers, hotels, attractions, etc.), you **MUST** include real latitude and longitude.
  Use accurate coordinates for real Tesla Superchargers (examples: London Battersea ≈ 51.477, -0.17; Birmingham NEC ≈ 52.45, -1.72; Leeds Skelton Lake ≈ 53.78, -1.46).
- The output must contain the itinerary in this **exact JSON shape** after your normal reply:

ITINERARY_UPDATE:
{
  "days": [
    {
      "day": 1,
      "stops": [
        {
          "id": "unique-string",
          "name": "Human readable name",
          "type": "supercharger" | "hotel" | "attraction" | "restaurant" | "custom",
          "lat": 51.477,
          "lng": -0.17,
          "day": 1,
          "order": 1,
          "estArrivalBattery": 25,
          "chargeMinutes": 25,
          "notes": "optional short note"
        }
      ]
    }
  ],
  "summary": {
    "totalDistanceKm": 650,
    "estDriveHours": 10.5,
    "estChargeHours": 1.5,
    "superchargers": 3,
    "hotels": 1
  }
}

Rules:
- Never mention the JSON in your spoken reply.
- Make realistic assumptions for a first draft (max 5-6h driving/day, include major Superchargers, logical overnights).
- After the draft, ask what the user wants to change.

Conversation:
${messages.map(m => `${m.role.toUpperCase()}: ${m.content}`).join('\n')}

ASSISTANT:`;
  }

  async chat(messages: ChatMessage[], itinerary: any, vehicle: string): Promise<GrokResponse> {
    const requestId = crypto.randomUUID().slice(0, 8);

    log.info({ requestId, vehicle: vehicle.name, messageCount: messages.length }, '=== NEW CHAT REQUEST ===');

    if (env.xaiApiKey) {
      log.info({ requestId }, 'Using real xAI API');
      return this.callXaiApi(messages, itinerary, vehicle, requestId);
    }

    if (this.useFallback) {
      return this.dumbFallback(messages, requestId);
    }

    const prompt = this.buildPrompt(messages, itinerary, vehicle);
    const tmp = await mkdtemp(join(tmpdir(), 'grok-eu-'));

    try {
      const args = ['-p', prompt, '--output-format', 'json', '--yolo', '--disallowed-tools', 'run_terminal_cmd,search_replace,write_file,Agent', '--tools', 'web_search,web_fetch', '--max-turns', '6', '--effort', 'high', '--cwd', tmp];

      const result = await new Promise<string>((resolve, reject) => {
        const child = spawn(env.grokBin, args, { cwd: tmp, env: { ...process.env } });
        let stdout = '';
        child.stdout.on('data', d => stdout += d);
        child.on('close', code => code === 0 ? resolve(stdout) : reject(new Error(`grok exited ${code}`)));
      });

      const data = JSON.parse(result);
      const rawText = data.text || '';

      const { cleanText, itinerary: parsed } = this.extractItineraryUpdate(rawText);

      return { text: cleanText, updatedItinerary: parsed };

    } catch (err) {
      log.error({ requestId, err }, 'Local grok CLI failed — falling back to xAI API');
      if (env.xaiApiKey) {
        return this.callXaiApi(messages, itinerary, vehicle, requestId);
      }
      return this.dumbFallback(messages, requestId);
    } finally {
      await rm(tmp, { recursive: true, force: true }).catch(() => {});
    }
  }

  private async callXaiApi(messages: ChatMessage[], itinerary: any, vehicle: string, requestId: string): Promise<GrokResponse> {
    const prompt = this.buildPrompt(messages, itinerary, vehicle);

    log.info({ requestId, promptLength: prompt.length }, 'Calling real xAI API');

    try {
      const response = await fetch('https://api.x.ai/v1/chat/completions', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
          'Authorization': `Bearer ${env.xaiApiKey}`,
        },
        body: JSON.stringify({
          model: 'grok-3',
          messages: [{ role: 'user', content: prompt }],
          temperature: 0.7,
        }),
      });

      if (!response.ok) {
        const text = await response.text();
        log.error({ requestId, status: response.status }, 'xAI API error');
        return this.dumbFallback(messages, requestId);
      }

      const data = await response.json();
      const rawText = data.choices?.[0]?.message?.content || '';

      const { cleanText, itinerary: parsed } = this.extractItineraryUpdate(rawText);

      return { text: cleanText, updatedItinerary: parsed };
    } catch (err) {
      log.error({ requestId, err }, 'xAI API call failed');
      return this.dumbFallback(messages, requestId);
    }
  }

  private extractItineraryUpdate(text: string): { cleanText: string; itinerary: any | null } {
    const upperText = text.toUpperCase();
    const sentinelIndex = upperText.indexOf(SENTINEL.toUpperCase());

    if (sentinelIndex === -1) {
      return { cleanText: text.trim(), itinerary: null };
    }

    const after = text.substring(sentinelIndex + SENTINEL.length).trim();

    // Try to find a JSON object, even if wrapped in ```json
    let jsonStart = after.indexOf('{');
    if (jsonStart === -1) return { cleanText: text.substring(0, sentinelIndex).trim(), itinerary: null };

    // Find matching closing brace
    let depth = 0;
    let end = -1;
    let inString = false;
    let escape = false;

    for (let i = jsonStart; i < after.length; i++) {
      const ch = after[i];
      if (escape) { escape = false; continue; }
      if (ch === '\\') { escape = true; continue; }
      if (ch === '"') { inString = !inString; continue; }
      if (!inString) {
        if (ch === '{') depth++;
        if (ch === '}') {
          depth--;
          if (depth === 0) { end = i; break; }
        }
      }
    }

    if (end === -1) return { cleanText: text.substring(0, sentinelIndex).trim(), itinerary: null };

    const jsonStr = after.substring(jsonStart, end + 1);

    try {
      const parsed = JSON.parse(jsonStr);
      const cleanText = text.substring(0, sentinelIndex).trim();
      return { cleanText, itinerary: parsed };
    } catch (e) {
      log.error({ err: e }, 'Failed to parse ITINERARY_UPDATE JSON');
      return { cleanText: text.substring(0, sentinelIndex).trim(), itinerary: null };
    }
  }

  private async dumbFallback(messages: ChatMessage[], requestId: string): Promise<GrokResponse> {
    const last = messages[messages.length - 1]?.content.toLowerCase() || '';
    if (['hi', 'hello', 'hey'].some(g => last.includes(g))) {
      return { text: "Hello! I'm Grok Drive. How can I help plan your UK or European Tesla trip today?", updatedItinerary: null };
    }
    return {
      text: "I'm ready to plan a great Tesla route for you across the UK and Europe. Tell me where you want to go!",
      updatedItinerary: null
    };
  }
}

export const grok = new GrokHeadlessClient();