import axios, { AxiosInstance } from 'axios';
import sharp from 'sharp';
import path from 'path';
import { randomUUID } from 'crypto';
import { z } from 'zod';
import { DEFAULT_OUTPUT_DIR, CF_IMGBED_UPLOAD_URL, CF_IMGBED_API_KEY } from '../config.js';
import { ensureDirectoryExists } from '../utils/fileUtils.js';
import { uploadToCfImgbed } from '../utils/cfUtils.js';

// --- Interfaces ---
export interface GenerateImageArgs {
  prompt: string;
  negative_prompt?: string;
  steps?: number;
  width?: number;
  height?: number;
  cfg_scale?: number;
  sampler_name?: string;
  scheduler_name?: string;
  seed?: number;
  batch_size?: number;
  restore_faces?: boolean;
  tiling?: boolean;
  output_path?: string;
  distilled_cfg_scale?: number;
}

interface SDAPIPayload {
  prompt: string;
  negative_prompt: string;
  steps: number;
  width: number;
  height: number;
  cfg_scale: number;
  sampler_name: string;
  scheduler_name: string;
  seed: number;
  n_iter: number; // Corresponds to batch_size
  restore_faces?: boolean;
  tiling?: boolean;
  distilled_cfg_scale?: number;
}

// --- Zod Schema for Input Validation ---
export const generateImageSchema = z.object({
  prompt: z.string().min(1, "Prompt cannot be empty"),
  negative_prompt: z.string().optional(),
  steps: z.number().int().min(1).max(150).optional().default(20),
  width: z.number().int().min(512).max(2048).optional().default(1024),
  height: z.number().int().min(512).max(2048).optional().default(1024),
  cfg_scale: z.number().min(1).max(30).optional().default(3.5),
  sampler_name: z.string().optional().default('Euler a'),
  scheduler_name: z.string().optional().default('Automatic'),
  seed: z.number().int().optional().default(-1),
  batch_size: z.number().int().min(1).max(4).optional().default(1),
  restore_faces: z.boolean().optional().default(false),
  tiling: z.boolean().optional().default(false),
  output_path: z.string().optional(),
  distilled_cfg_scale: z.number().min(1).max(30).optional().default(3.5),
});

// Type alias for validated arguments
export type ValidatedGenerateImageArgs = z.infer<typeof generateImageSchema>;


// --- Helper Functions ---

// Removed ensureDirectoryExists as it will be imported
// Removed uploadToCfImgbed as it will be imported

function sanitizePromptForFilename(prompt: string, maxLength: number = 50): string {
  if (!prompt) {
    return 'no_prompt';
  }
  // Truncate the prompt to maxLength before further processing
  let sanitized = prompt.substring(0, maxLength);

  sanitized = sanitized.replace(/[^\p{L}\p{N}\s_-]/gu, '').replace(/\s+/g, '_');
  // Replace multiple underscores with a single one
  sanitized = sanitized.replace(/__+/g, '_');
  // Remove leading/trailing underscores
  sanitized = sanitized.replace(/^_+|_+$/g, '');

  return sanitized.toLowerCase() || 'prompt'; // Ensure not empty
}

// --- Main Handler Function ---
export async function handleGenerateImage(
  args: ValidatedGenerateImageArgs,
  axiosInstance: AxiosInstance
): Promise<{ content: { type: string; text: string }[] }> {
  const outputDir = args.output_path ? path.normalize(args.output_path.trim()) : DEFAULT_OUTPUT_DIR;
  await ensureDirectoryExists(outputDir);

  const payload: SDAPIPayload = {
    prompt: args.prompt,
    negative_prompt: args.negative_prompt || 'worst quality,bad quality,bad hands,very displeasing,extra digit,fewer digits,jpeg artifacts,signature,username,reference,mutated,lineup,manga,comic,disembodied,futanari,yaoi,dickgirl,turnaround,2koma,4koma,monster,cropped,amputee,text,bad foreshortening,what,guro,logo,bad anatomy,bad perspective,bad proportions,artistic error,anatomical nonsense,amateur,out of frame,multiple views,disfigured,ugly,mutation,acnes,skin spots,skin blemishes,poorly drawn face,bathtub,shore,grass,buildings,stone,missing fingers,fused fingers,disconnected limbs,extra limb,extra arms,mutated hands,poorly drawn hands,malformed hands,mutated hands and fingers,missing limb,malformed limbs,deformed',
    steps: args.steps,
    width: args.width,
    height: args.height,
    cfg_scale: args.cfg_scale,
    sampler_name: args.sampler_name,
    seed: args.seed,
    n_iter: args.batch_size, // Maps to n_iter in API
    distilled_cfg_scale: args.distilled_cfg_scale,
    scheduler_name: args.scheduler_name,
    tiling: args.tiling,
    restore_faces: args.restore_faces,
  };

  const response = await axiosInstance.post('/sdapi/v1/txt2img', payload);
  if (!response.data.images?.length) {
    throw new Error('No images generated by SD API');
  }

  const results = [];
  // For local filename, use maxLength 50 (current default of sanitizePromptForFilename)
  const localSanitizedPrompt = sanitizePromptForFilename(args.prompt); 
  // For remote filename for ImgBed, use a much larger maxLength, e.g., 1000 (or a configurable value)
  const remoteSanitizedPrompt = sanitizePromptForFilename(args.prompt, 1000); 

  for (const imageData of response.data.images) {
    const base64Data = imageData.includes(',') ? imageData.split(',')[1] : imageData;
    
    // Filename for local storage (truncated prompt)
    const localFilename = `sd_${localSanitizedPrompt}_${randomUUID()}.png`;
    const outputPath = path.join(outputDir, localFilename);
    const imageBuffer = Buffer.from(base64Data, 'base64');

    await sharp(imageBuffer).toFile(outputPath);

    let uploadedUrl: string | null = null;
    if (CF_IMGBED_UPLOAD_URL && CF_IMGBED_API_KEY) {
      // Filename for ImgBed upload (longer, sanitized prompt)
      const remoteFilename = `sd_${remoteSanitizedPrompt}_${randomUUID()}.png`;
      uploadedUrl = await uploadToCfImgbed(imageBuffer, remoteFilename);
    }

    results.push({
      path: outputPath,
      url: uploadedUrl,
    });
  }

  return { content: [{ type: 'text', text: JSON.stringify(results) }] };
}