import fs from 'fs/promises';
import path from 'path';

const OLLAMA_API = 'http://127.0.0.1:11434/api/generate';
const MODEL = "llava:latest";

const SYSTEM_PROMPT = `Du bist ein KI-Assistent zur Klassifizierung von Bildern für ein Infrastruktur-Unternehmen. 
Bitte analysiere das Bild und antworte AUSSCHLIESSLICH mit einem validen JSON-Objekt, das folgendem Schema entspricht:
{
  "category": "pv" | "wind" | "fiber" | "power" | "battery" | "unknown",
  "environment": "freies_feld" | "urban" | "wald" | "indoor" | "unknown",
  "quality_score": number, // 1-10, wie gut eignet sich das Bild als Hero/Titelbild? (Schärfe, Ästhetik, Belichtung)
  "description": "string" // Kurze deutsche Beschreibung, max 10 Wörter
}

Erklärung der Kategorien:
- "pv": Photovoltaik-Anlagen, Solarparks, Solarmodule
- "wind": Windräder, Windparks
- "fiber": Kabeltiefbau, Spülbohrtechnik, Bagger bei Kabelverlegung, Glasfaser
- "power": Umspannwerke, Hochspannungstrassen, Strommasten
- "battery": Batterie-Speichersysteme (BESS), Container-Speicher

Erklärung der Umgebung:
- "freies_feld": Offenes Feld, Wiese, Natur (vom Kunden stark bevorzugt für PV und Wind!)
- "urban": Stadt, Straße, Siedlung
- "wald": Im Wald, viele Bäume nah dran
- "indoor": Innenräume

WICHTIG: Gib NUR JSON zurück. Keinen Markdown-Code-Block. Nichts anderes.
`;

interface ImageMeta {
  file: string;
  category: 'pv' | 'wind' | 'fiber' | 'power' | 'battery' | 'unknown';
  environment: 'freies_feld' | 'urban' | 'wald' | 'indoor' | 'unknown';
  quality_score: number;
  description: string;
}

async function classifyImage(filePath: string): Promise<ImageMeta | null> {
  try {
    const ext = path.extname(filePath).toLowerCase();
    if (!['.jpg', '.jpeg', '.png', '.webp', '.avif'].includes(ext)) {
      return null;
    }

    const imageBuffer = await fs.readFile(filePath);
    const base64Image = imageBuffer.toString('base64');

    console.log(`Analyzing ${path.basename(filePath)}...`);

    const response = await fetch(OLLAMA_API, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        model: MODEL,
        prompt: "Bitte klassifiziere dieses Bild gemäß den Anweisungen und antworte nur mit JSON.",
        system: SYSTEM_PROMPT,
        images: [base64Image],
        stream: false,
        format: "json",
        options: {
          temperature: 0.1,
          num_predict: 200,
        }
      })
    });

    if (!response.ok) {
      console.error(`Ollama error for ${filePath}: ${response.status} ${response.statusText}`);
      return null;
    }

    const result = await response.json();
    let responseText = result.response.trim();
    
    // Fallback if model wraps in code blocks
    if (responseText.startsWith('\`\`\`json')) {
      responseText = responseText.replace(/^\`\`\`json/m, '').replace(/\`\`\`$/m, '').trim();
    }

    const parsed = JSON.parse(responseText);
    
    return {
      file: '/assets/photos/' + path.basename(filePath),
      category: parsed.category || 'unknown',
      environment: parsed.environment || 'unknown',
      quality_score: parsed.quality_score || 5,
      description: parsed.description || ''
    };

  } catch (error) {
    console.error(`Error processing ${filePath}:`, error);
    return null;
  }
}

async function main() {
  const photosDir = path.join(process.cwd(), 'public', 'assets', 'photos');
  const files = await fs.readdir(photosDir);
  const metadataMap: Record<string, ImageMeta> = {};
  
  const metadataPath = path.join(process.cwd(), 'public', 'assets', 'image-metadata.json');
  
  // Load existing to skip if needed, or to resume
  let existing: Record<string, ImageMeta> = {};
  try {
    const existingData = await fs.readFile(metadataPath, 'utf-8');
    existing = JSON.parse(existingData);
  } catch (e) {
    // Ignore if not exists
  }

  for (const file of files) {
    const filePath = path.join(photosDir, file);
    const relativePath = '/assets/photos/' + file;
    
    if (existing[relativePath] && existing[relativePath].category !== 'unknown') {
      console.log(`Skipping ${file}, already analyzed.`);
      metadataMap[relativePath] = existing[relativePath];
      continue;
    }

    const meta = await classifyImage(filePath);
    if (meta) {
      metadataMap[relativePath] = meta;
      // Save incrementally so we don't lose progress if it crashes
      await fs.writeFile(metadataPath, JSON.stringify(metadataMap, null, 2));
    }
  }

  console.log('Classification complete!');
}

main().catch(console.error);