import PQueue from 'p-queue';
import sharp from 'sharp';
import { generateText } from 'ai';
import { createAnthropic } from '@ai-sdk/anthropic';
import { createOpenAI } from '@ai-sdk/openai';
import { createGoogleGenerativeAI } from '@ai-sdk/google';
import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
import { Job } from '../db/models.js';
import { broadcast } from '../ws/broadcast.js';
import { findModel, DEFAULT_MODEL_ID, normalizeModelId } from '../models.js';

// ---------------------------------------------------------------------------
// Provider instances
// ---------------------------------------------------------------------------
const anthropic = createAnthropic({
  apiKey: process.env.ANTHROPIC_API_KEY,
});

const openai = createOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
});

const google = createGoogleGenerativeAI({
  apiKey: process.env.GOOGLE_API_KEY,
});

// Ollama Cloud exposes an OpenAI-compatible /v1 endpoint.
// Using @ai-sdk/openai-compatible avoids the local-Ollama schema validation
// in ollama-ai-provider which requires fields (eval_duration etc.) that
// Ollama Cloud doesn't return.
const ollamaCloud = createOpenAICompatible({
  name:    'ollama-cloud',
  baseURL: 'https://ollama.com/v1',
  headers: {
    Authorization: `Bearer ${process.env.OLLAMA_API_KEY ?? ''}`,
  },
});

const PROVIDERS = {
  anthropic: (id) => anthropic(id),
  openai:    (id) => openai(id),
  google:    (id) => google(id),
  ollama:    (id) => ollamaCloud(id),
};

function resolveModel(modelId) {
  const normalized = normalizeModelId(modelId);
  const meta = findModel(normalized) ?? findModel(DEFAULT_MODEL_ID);
  return PROVIDERS[meta.provider](meta.id);
}

function resolveModelMeta(modelId) {
  const normalized = normalizeModelId(modelId);
  const meta = findModel(normalized) ?? findModel(DEFAULT_MODEL_ID);
  return { normalized, meta };
}

// ---------------------------------------------------------------------------
// p-queue: in-process queue, no external server needed
// ---------------------------------------------------------------------------
export const queue = new PQueue({
  concurrency: parseInt(process.env.JOB_CONCURRENCY ?? '3', 10),
});

queue.on('add',  () => broadcastQueueStats());
queue.on('active', () => broadcastQueueStats());
queue.on('next', () => broadcastQueueStats());
queue.on('idle', () => broadcastQueueStats());

function broadcastQueueStats() {
  broadcast({ type: 'queue_stats', pending: queue.size, running: queue.pending });
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
async function setStatus(job, status, extra = {}) {
  await job.update({ status, ...extra });
  broadcast({ type: 'job_update', job: job.toJSON() });
}

/** Padding factor applied to each side of the detected bounding box. */
const BBOX_PADDING = parseFloat(process.env.BBOX_PADDING ?? '0.15');

/**
 * Extract the raw image buffer + mime type from a data-URL string.
 */
function parseDataUrl(dataUrl) {
  const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/s);
  if (!match) throw new Error('Invalid data URL');
  return { mimeType: match[1], buffer: Buffer.from(match[2], 'base64') };
}

/**
 * Build a data-URL from a buffer + mime type.
 */
function toDataUrl(buffer, mimeType) {
  return `data:${mimeType};base64,${buffer.toString('base64')}`;
}

// ---------------------------------------------------------------------------
// Pass 1 — Ask the LLM to locate the subject described in the prompt
// ---------------------------------------------------------------------------
const BBOX_SYSTEM_PROMPT = `You are a vision assistant that locates objects in images.
The user will give you an image and a description of what they are interested in.
Your job is to identify the main subject described and return the bounding box
coordinates for that subject as a proportion of the image dimensions (values 0-1).

You MUST respond with ONLY a JSON object in this exact format, no other text:
{
  "subject": "<short description of what you identified>",
  "bbox": {
    "x": <left edge 0-1>,
    "y": <top edge 0-1>,
    "width": <box width 0-1>,
    "height": <box height 0-1>
  }
}

If the prompt doesn't clearly reference a specific subject in the image, or if the
subject fills most of the image already, return the full image:
{ "subject": "full image", "bbox": { "x": 0, "y": 0, "width": 1, "height": 1 } }`;

async function detectSubjectBbox(model, imageDataUrl, prompt) {
  const { text } = await generateText({
    model,
    system: BBOX_SYSTEM_PROMPT,
    messages: [
      {
        role: 'user',
        content: [
          {
            type: 'text',
            text: `Locate the main subject referenced in this prompt and return its bounding box.\n\nPrompt: "${prompt}"`,
          },
          { type: 'image', image: imageDataUrl },
        ],
      },
    ],
    maxTokens: 256,
  });

  // Parse JSON from the response — tolerate markdown fences
  const cleaned = text.replace(/```json\s*/g, '').replace(/```/g, '').trim();
  const parsed = JSON.parse(cleaned);

  // Validate and clamp values
  const bbox = parsed.bbox;
  if (
    typeof bbox?.x !== 'number' ||
    typeof bbox?.y !== 'number' ||
    typeof bbox?.width !== 'number' ||
    typeof bbox?.height !== 'number'
  ) {
    throw new Error('LLM returned invalid bbox structure');
  }

  const clamp = (v) => Math.max(0, Math.min(1, v));
  return {
    subject: parsed.subject ?? 'unknown',
    bbox: {
      x:      clamp(bbox.x),
      y:      clamp(bbox.y),
      width:  clamp(bbox.width),
      height: clamp(bbox.height),
    },
  };
}

// ---------------------------------------------------------------------------
// Crop helper — takes a data-URL, bounding box (0-1), and padding factor
// ---------------------------------------------------------------------------
async function cropImage(imageDataUrl, bbox, padding = BBOX_PADDING) {
  const { mimeType, buffer } = parseDataUrl(imageDataUrl);
  const metadata = await sharp(buffer).metadata();
  const imgW = metadata.width;
  const imgH = metadata.height;

  // Convert proportional bbox to pixel values
  let left   = bbox.x * imgW;
  let top    = bbox.y * imgH;
  let width  = bbox.width * imgW;
  let height = bbox.height * imgH;

  // Add padding (proportional to the box dimensions)
  const padX = width * padding;
  const padY = height * padding;
  left   = Math.max(0, left - padX);
  top    = Math.max(0, top - padY);
  width  = Math.min(imgW - left, width + padX * 2);
  height = Math.min(imgH - top, height + padY * 2);

  // Round to integers for sharp
  left   = Math.round(left);
  top    = Math.round(top);
  width  = Math.round(width);
  height = Math.round(height);

  // Guard against degenerate crops
  if (width < 1 || height < 1) {
    return { croppedDataUrl: imageDataUrl, cropPixels: null };
  }

  const croppedBuf = await sharp(buffer)
    .extract({ left, top, width, height })
    .toBuffer();

  return {
    croppedDataUrl: toDataUrl(croppedBuf, mimeType),
    cropPixels: { left, top, width, height, imgW, imgH },
  };
}

// ---------------------------------------------------------------------------
// Main runner — two-pass: detect bbox → crop → analyse cropped image
// ---------------------------------------------------------------------------
async function runJob(jobId) {
  const job = await Job.findByPk(jobId);
  if (!job) return;

  await setStatus(job, 'running');

  try {
    resolveModelMeta(job.model);
    const model = resolveModel(job.model);

    // ---- Pass 1: detect subject bounding box ----
    let bboxResult;
    try {
      bboxResult = await detectSubjectBbox(model, job.imageDataUrl, job.prompt);
    } catch (bboxErr) {
      // If bbox detection fails, fall back to the full image
      console.warn(`Bbox detection failed, using full image: ${bboxErr.message}`);
      bboxResult = {
        subject: 'full image (fallback)',
        bbox: { x: 0, y: 0, width: 1, height: 1 },
      };
    }

    // Persist the detected bbox for debugging / UI display
    await job.update({
      detectedSubject: bboxResult.subject,
      bboxX:      bboxResult.bbox.x,
      bboxY:      bboxResult.bbox.y,
      bboxWidth:  bboxResult.bbox.width,
      bboxHeight: bboxResult.bbox.height,
    });
    broadcast({ type: 'job_update', job: job.toJSON() });

    // ---- Crop the image around the detected subject ----
    const isFullImage =
      bboxResult.bbox.x === 0 &&
      bboxResult.bbox.y === 0 &&
      bboxResult.bbox.width === 1 &&
      bboxResult.bbox.height === 1;

    let imageForAnalysis = job.imageDataUrl;

    if (!isFullImage) {
      const { croppedDataUrl, cropPixels } = await cropImage(
        job.imageDataUrl,
        bboxResult.bbox,
      );
      imageForAnalysis = croppedDataUrl;

      // Optionally store the cropped image for UI preview
      await job.update({ croppedImageDataUrl: croppedDataUrl });
      broadcast({ type: 'job_update', job: job.toJSON() });

      if (cropPixels) {
        console.log(
          `Cropped to ${cropPixels.width}×${cropPixels.height} ` +
          `from ${cropPixels.imgW}×${cropPixels.imgH} ` +
          `(subject: "${bboxResult.subject}")`,
        );
      }
    }

    // ---- Pass 2: run the original prompt against the (cropped) image ----
    const { text, usage } = await generateText({
      model,
      messages: [
        {
          role: 'user',
          content: [
            { type: 'text',  text: job.prompt },
            { type: 'image', image: imageForAnalysis },
          ],
        },
      ],
      maxTokens: 1024,
    });

    await setStatus(job, 'done', {
      result:       text,
      inputTokens:  usage?.promptTokens     ?? null,
      outputTokens: usage?.completionTokens ?? null,
    });
  } catch (err) {
    const detail = err?.message
      ? `${err.name ?? 'Error'}: ${err.message}${err.cause ? `\nCause: ${JSON.stringify(err.cause, null, 2)}` : ''}`
      : String(err);
    await setStatus(job, 'error', { errorMessage: detail });
  }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
export function enqueueJob(jobId) {
  queue.add(() => runJob(jobId));
}