import { NextResponse } from 'next/server'
import { aiRatelimit, getClientIp, checkRateLimit } from '@/lib/ratelimit'
import { sanitizePrompt } from '@/lib/validation'
import { getArchetypeById } from '@/data/archetypes'
import {
  buildPrimePrompt,
  DEFAULT_AURA_COLOR,
  PRIME_TIERS,
  PRIME_GRAINS,
  PRIME_RAINS,
  type PrimeTier,
  type PrimeGrain,
  type PrimeRain,
} from '@/lib/primePrompt'

import Replicate from 'replicate'

// Replicate SDK requires Node runtime (Edge does not support it reliably).
export const runtime = 'nodejs'
export const dynamic = 'force-dynamic'

const MAX_FILE_MB = 10
const MAX_FILE_BYTES = MAX_FILE_MB * 1024 * 1024

const replicateApiToken = process.env.REPLICATE_API_TOKEN
const replicateClient = replicateApiToken
  ? new Replicate({ auth: replicateApiToken })
  : null

// SDXL for img2img style transfer with identity preservation
const SDXL_MODEL =
  'stability-ai/sdxl:7762fd07cf82c948538e41f63f77d685e02b063e37e496e96eefd46c929f9bdc'

function parseDataUri(dataUri: string): { mime: string; buffer: Buffer } | null {
  const match = dataUri.match(/^data:(image\/[a-zA-Z0-9+.-]+);base64,(.+)$/)
  if (!match) return null
  return { mime: match[1].toLowerCase(), buffer: Buffer.from(match[2], 'base64') }
}

function normalizeAuraColor(input: string | null): string {
  if (!input) return DEFAULT_AURA_COLOR
  const cleaned = sanitizePrompt(input)
    .toLowerCase()
    .replace(/[^a-z0-9\s-]/g, '')
    .trim()
  return cleaned || DEFAULT_AURA_COLOR
}

function normalizeTier(input: string | null): PrimeTier | null {
  if (!input) return 'core'
  const value = input.toLowerCase().trim()
  return PRIME_TIERS.includes(value as PrimeTier) ? (value as PrimeTier) : null
}

function normalizeGrain(input: string | null): PrimeGrain | null {
  if (!input) return 'gritty'
  const value = input.toLowerCase().trim()
  return PRIME_GRAINS.includes(value as PrimeGrain) ? (value as PrimeGrain) : null
}

function normalizeRain(input: string | null): PrimeRain | null {
  if (!input) return 'light'
  const value = input.toLowerCase().trim()
  return PRIME_RAINS.includes(value as PrimeRain) ? (value as PrimeRain) : null
}

// GET - Poll Prediction Status
export async function GET(request: Request) {
  if (!replicateClient) {
    return NextResponse.json(
      { error: 'Replicate API token not configured' },
      { status: 500 }
    )
  }

  const url = new URL(request.url)
  const predictionId = url.searchParams.get('id')
  if (!predictionId) {
    return NextResponse.json(
      { error: 'Prediction id is required' },
      { status: 400 }
    )
  }

  try {
    const prediction = await replicateClient.predictions.get(predictionId)
    console.log('Replicate prediction response:', prediction)
    return NextResponse.json({
      id: prediction.id,
      status: prediction.status,
      output: prediction.output,
      error: prediction.error,
    })
  } catch (error: unknown) {
    console.error('Prime status check error:', error)
    const message = error instanceof Error ? error.message : 'Failed to fetch status'
    return NextResponse.json({ error: message }, { status: 500 })
  }
}

// POST - Start Prediction
export async function POST(request: Request) {
  try {
    if (!replicateClient) {
      return NextResponse.json(
        { error: 'Replicate API token not configured' },
        { status: 500 }
      )
    }
    const ip = getClientIp(request)
    const { success, error } = await checkRateLimit(aiRatelimit, ip)
    if (!success) return error

    // Parse JSON body
    let body
    try {
      body = await request.json()
    } catch {
      return NextResponse.json({ error: 'Invalid JSON body' }, { status: 400 })
    }

    const { image, tier, grain, rain, aura_color, archetype } = body

    if (!image || typeof image !== 'string' || !image.startsWith('data:image/')) {
      return NextResponse.json({ error: 'Invalid image data' }, { status: 400 })
    }

    const parsed = parseDataUri(image)
    if (!parsed) {
      return NextResponse.json({ error: 'Invalid image data' }, { status: 400 })
    }
    if (!['image/png', 'image/jpeg', 'image/webp', 'image/jpg'].includes(parsed.mime)) {
      return NextResponse.json({ error: 'Unsupported image format' }, { status: 400 })
    }
    if (parsed.buffer.length > MAX_FILE_BYTES) {
      return NextResponse.json({ error: 'Image too large' }, { status: 400 })
    }

    const normalizedTier = normalizeTier(tier)
    const normalizedGrain = normalizeGrain(grain)
    const normalizedRain = normalizeRain(rain)

    if (!normalizedTier || !normalizedGrain || !normalizedRain) {
      return NextResponse.json({ error: 'Invalid parameters' }, { status: 400 })
    }

    const auraColor = normalizeAuraColor(aura_color)
    const archetypeId = archetype?.toString().toLowerCase().trim()
    const archetypeData = archetypeId ? getArchetypeById(archetypeId) : undefined

    // Build the PRIME prompt with proper structure
    const { prompt, negativePrompt } = buildPrimePrompt({
      tier: normalizedTier,
      auraColor,
      archetypeName: archetypeData?.name,
      archetypeModifier: archetypeData?.promptModifier,
      grain: normalizedGrain,
      rain: normalizedRain,
    })

    // Create prediction with SDXL img2img
    // Using original working parameters to avoid timeout issues
    const prediction = await replicateClient.predictions.create({
      version: SDXL_MODEL.split(':')[1],
      input: {
        image,
        prompt,
        negative_prompt: negativePrompt,
        width: 1024,
        height: 1024,
        num_inference_steps: 25, // Keep reduced for speed
        guidance_scale: 7.5,
        prompt_strength: 0.8, // Increase for more noticeable transformation
        scheduler: 'K_EULER', // Revert to original
        num_outputs: 1,
        apply_watermark: false,
      },
    })

    return NextResponse.json(
      { id: prediction.id, status: prediction.status },
      { status: 201 }
    )
  } catch (error: unknown) {
    console.error('Prime transform error:', error)
    const message = error instanceof Error ? error.message : 'Failed to transform image'
    return NextResponse.json({ error: message }, { status: 500 })
  }
}