import { S3Client, GetObjectCommand } from "@aws-sdk/client-s3";
import { sql } from "@/db";
import { logAudit } from "@/lib/audit";
import OpenAI from "openai";

const VISION_MODEL = process.env.VISION_MODEL || "gemini-flash";
const EMBEDDING_MODEL = process.env.EMBEDDING_MODEL || "text-embedding-3-small";
const LITELLM_URL = process.env.LITELLM_BASE_URL;
const LITELLM_KEY = process.env.LITELLM_API_KEY;

function getR2Config() {
  return {
    accountId: process.env.R2_ACCOUNT_ID!,
    accessKeyId: process.env.R2_ACCESS_KEY_ID!,
    secretKey: process.env.R2_SECRET_ACCESS_KEY!,
    bucket: process.env.R2_BUCKET_NAME!,
    publicUrl: process.env.R2_PUBLIC_URL,
  };
}

async function downloadFromR2(key: string): Promise<Buffer> {
  const R2 = getR2Config();
  const client = new S3Client({
    region: "auto",
    endpoint: `https://${R2.accountId}.r2.cloudflarestorage.com`,
    credentials: { accessKeyId: R2.accessKeyId, secretAccessKey: R2.secretKey },
  });
  const response = await client.send(new GetObjectCommand({ Bucket: R2.bucket, Key: key }));
  const chunks: Buffer[] = [];
  for await (const chunk of response.Body as AsyncIterable<Buffer>) {
    chunks.push(chunk);
  }
  return Buffer.concat(chunks);
}

interface VisionResult {
  caption: string;
  tags: string[];
}

async function getVisionCaption(imageBuffer: Buffer): Promise<VisionResult> {
  if (!LITELLM_URL || !LITELLM_KEY) {
    throw new Error("LiteLLM not configured");
  }

  const base64 = imageBuffer.toString("base64");
  const response = await fetch(`${LITELLM_URL}/v1/chat/completions`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "Authorization": `Bearer ${LITELLM_KEY}`,
    },
    body: JSON.stringify({
      model: VISION_MODEL,
      messages: [
        {
          role: "user",
          content: [
            {
              type: "text",
              text: `Describe this photo of a baby or child in 2-3 sentences. Focus on what's happening, who's in it, mood, and location if obvious. Then list 5-10 descriptive tags. Respond ONLY with valid JSON in this exact format: {"caption": "...", "tags": ["...", "..."]}`,
            },
            {
              type: "image_url",
              image_url: { url: `data:image/webp;base64,${base64}` },
            },
          ],
        },
      ],
      max_tokens: 300,
      temperature: 0.2,
      response_format: { type: "json_object" },
    }),
  });

  if (!response.ok) {
    throw new Error(`Vision API error: ${response.status}`);
  }

  const data = await response.json();
  const content = data.choices?.[0]?.message?.content;
  if (!content) throw new Error("Empty vision response");

  return JSON.parse(content) as VisionResult;
}

async function getEmbedding(text: string): Promise<number[]> {
  const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY || LITELLM_KEY, baseURL: process.env.OPENAI_API_KEY ? undefined : `${LITELLM_URL}/v1` });
  const res = await openai.embeddings.create({ model: EMBEDDING_MODEL, input: text });
  return res.data[0].embedding;
}

export async function processMemoryVision(memoryId: string): Promise<void> {
  const rows = await sql`SELECT * FROM memories WHERE id = ${memoryId} LIMIT 1`;
  const memory = rows[0];
  if (!memory) throw new Error(`Memory not found: ${memoryId}`);

  // Skip vision for private memories
  if (memory.is_private) {
    await sql`UPDATE memories SET processing_status = 'ready', updated_at = now() WHERE id = ${memoryId}`;
    return;
  }

  try {
    // Use thumbnail if available (smaller payload), else original
    const key = memory.r2_thumbnail_key || memory.r2_key;
    const imageBuffer = await downloadFromR2(key);

    // Get caption + tags from vision model
    const { caption, tags } = await getVisionCaption(imageBuffer);

    // Generate embedding from caption
    const embedding = await getEmbedding(caption);
    // Format as pgvector literal: [0.1,0.2,...]
    const embeddingLiteral = `[${embedding.join(",")}]`;

    await sql`
      UPDATE memories SET
        vision_caption = ${caption},
        vision_tags = ${tags},
        vision_embedding = ${embeddingLiteral}::vector,
        processing_status = 'ready',
        updated_at = now()
      WHERE id = ${memoryId}
    `;
  } catch (e) {
    console.error(`[vision] Failed for memory ${memoryId}:`, e);
    // Vision is optional — mark as ready so the image is still visible.
    // Only the confirm step (HeadObject failure) should mark as 'failed'.
    await sql`UPDATE memories SET processing_status = 'ready', updated_at = now() WHERE id = ${memoryId}`;
    await logAudit({
      action: "vision_processing_failed",
      metadata: { memoryId, error: String(e) },
      userId: memory.uploaded_by,
      familyId: memory.family_id,
    }).catch(() => {});
  }
}
