import { S3Client, GetObjectCommand } from "@aws-sdk/client-s3";
import { GARMENT_CATEGORIES, type GarmentCategory } from "@/db/schema/wardrobe";

const VISION_MODEL = process.env.VISION_MODEL || "gemini-flash";
const LITELLM_URL = process.env.LITELLM_BASE_URL;
const LITELLM_KEY = process.env.LITELLM_API_KEY;

interface GarmentVisionResult {
  name: string;
  category: GarmentCategory;
  colors: string[];
  seasons: string[];
  occasion_tags: string[];
}

const SAFE_DEFAULT: GarmentVisionResult = {
  name: "",
  category: "top",
  colors: [],
  seasons: [],
  occasion_tags: [],
};

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!,
  };
}

async function fetchFromR2(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 res = await client.send(new GetObjectCommand({ Bucket: R2.bucket, Key: key }));
  const chunks: Buffer[] = [];
  for await (const chunk of res.Body as AsyncIterable<Buffer>) {
    chunks.push(chunk);
  }
  return Buffer.concat(chunks);
}

export async function tagGarment(imageKey: string): Promise<GarmentVisionResult> {
  if (!LITELLM_URL || !LITELLM_KEY) {
    console.warn("[wardrobe/vision] LiteLLM not configured — returning defaults");
    return { ...SAFE_DEFAULT };
  }

  let imageBuffer: Buffer;
  try {
    imageBuffer = await fetchFromR2(imageKey);
  } catch (err) {
    console.error("[wardrobe/vision] R2 fetch failed:", err);
    return { ...SAFE_DEFAULT };
  }

  const base64 = imageBuffer.toString("base64");

  const systemPrompt = `You are a baby clothing tagger. Respond ONLY with a JSON object — no prose, no markdown fences.
The JSON must have exactly these keys:
- name: a short descriptive name for the garment (string, e.g. "Striped blue onesie")
- category: exactly one of: onesie|top|bottom|dress|outerwear|sleepwear|accessory
- colors: array of color names visible in the garment (e.g. ["blue","white"])
- seasons: subset of: summer|monsoon|winter (based on fabric weight and style)
- occasion_tags: subset of: everyday|daycare|festive|photoshoot
Do NOT attempt to guess the garment size. Do NOT include any other keys.`;

  let raw = "";
  try {
    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: "system",
            content: systemPrompt,
          },
          {
            role: "user",
            content: [
              {
                type: "image_url",
                image_url: { url: `data:image/jpeg;base64,${base64}` },
              },
              {
                type: "text",
                text: "Tag this baby garment.",
              },
            ],
          },
        ],
        max_tokens: 200,
        temperature: 0.1,
      }),
    });

    if (!response.ok) {
      console.error("[wardrobe/vision] API error:", response.status);
      return { ...SAFE_DEFAULT };
    }

    const data = await response.json();
    raw = data.choices?.[0]?.message?.content ?? "";
  } catch (err) {
    console.error("[wardrobe/vision] fetch failed:", err);
    return { ...SAFE_DEFAULT };
  }

  // Strip accidental markdown fences
  const cleaned = raw.replace(/^```(?:json)?\s*/i, "").replace(/\s*```\s*$/, "").trim();

  let parsed: Record<string, unknown>;
  try {
    parsed = JSON.parse(cleaned);
  } catch {
    console.warn("[wardrobe/vision] parse failed, raw:", raw);
    return { ...SAFE_DEFAULT };
  }

  // Validate category against allowed enum — treat hallucinated values as default
  const rawCategory = String(parsed.category ?? "").toLowerCase() as GarmentCategory;
  const safeCategory: GarmentCategory = GARMENT_CATEGORIES.includes(rawCategory)
    ? rawCategory
    : "top";

  const toStringArray = (v: unknown): string[] => {
    if (!Array.isArray(v)) return [];
    return v.filter((x): x is string => typeof x === "string");
  };

  return {
    name: typeof parsed.name === "string" ? parsed.name : "",
    category: safeCategory,
    colors: toStringArray(parsed.colors),
    seasons: toStringArray(parsed.seasons).filter(s =>
      ["summer", "monsoon", "winter"].includes(s)
    ),
    occasion_tags: toStringArray(parsed.occasion_tags).filter(t =>
      ["everyday", "daycare", "festive", "photoshoot"].includes(t)
    ),
  };
}
