import 'dotenv/config';
import express from 'express';
import multer from 'multer';
import { GoogleGenAI } from '@google/genai';

type InlineInput = {
  name?: string;
  mimeType?: string;
  data?: string;
  base64?: string;
  dataUrl?: string;
  url?: string;
};

type GenerateRequest = {
  prompt?: string;
  instruction?: string;
  model?: string;
  imageSize?: '1K' | '2K' | '4K';
  aspectRatio?: '1:1' | '4:3' | '3:4' | '16:9' | '9:16';
  images?: InlineInput[];
  documents?: InlineInput[];
  files?: InlineInput[];
};

type GeminiPart = {
  text?: string;
  inlineData?: {
    data: string;
    mimeType: string;
  };
};

const app = express();
const upload = multer({
  storage: multer.memoryStorage(),
  limits: {
    fileSize: Number(process.env.API_MAX_FILE_MB || 25) * 1024 * 1024,
    files: Number(process.env.API_MAX_FILES || 20),
  },
});

const apiKey = process.env.GEMINI_API_KEY || process.env.API_KEY;
const apiPort = Number(process.env.API_PORT || 3002);
const defaultImageModel = process.env.GEMINI_IMAGE_MODEL || 'gemini-3.1-flash-image-preview';
const defaultTextModel = process.env.GEMINI_TEXT_MODEL || 'gemini-2.5-flash';
const apiAuthToken = process.env.API_AUTH_TOKEN || '';

if (!apiKey) {
  console.warn('GEMINI_API_KEY/API_KEY is not set. API calls will fail until it is configured.');
}

const ai = new GoogleGenAI({ apiKey });

app.use((req, res, next) => {
  res.setHeader('Access-Control-Allow-Origin', process.env.API_CORS_ORIGIN || '*');
  res.setHeader('Access-Control-Allow-Methods', 'GET,POST,OPTIONS');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type,Authorization,x-api-key');
  if (req.method === 'OPTIONS') {
    res.status(204).end();
    return;
  }
  next();
});

app.use(express.json({ limit: process.env.API_JSON_LIMIT || '50mb' }));

function requireAuth(req: express.Request, res: express.Response, next: express.NextFunction) {
  if (!apiAuthToken) {
    next();
    return;
  }

  const authHeader = req.header('authorization') || '';
  const bearerToken = authHeader.startsWith('Bearer ') ? authHeader.slice(7) : '';
  const headerToken = req.header('x-api-key') || '';

  if (bearerToken === apiAuthToken || headerToken === apiAuthToken) {
    next();
    return;
  }

  res.status(401).json({
    ok: false,
    error: 'Unauthorized. Send Authorization: Bearer <API_AUTH_TOKEN> or x-api-key.',
  });
}

function stripDataPrefix(value: string) {
  const match = value.match(/^data:([^;]+);base64,(.+)$/);
  if (!match) {
    return null;
  }
  return {
    mimeType: match[1],
    data: match[2],
  };
}

async function fetchAsBase64(url: string) {
  const response = await fetch(url);
  if (!response.ok) {
    throw new Error(`Failed to fetch ${url}: ${response.status} ${response.statusText}`);
  }

  const contentLength = Number(response.headers.get('content-length') || 0);
  const maxBytes = Number(process.env.API_MAX_URL_MB || 25) * 1024 * 1024;
  if (contentLength > maxBytes) {
    throw new Error(`Remote file is too large. Max size is ${process.env.API_MAX_URL_MB || 25}MB.`);
  }

  const buffer = Buffer.from(await response.arrayBuffer());
  if (buffer.length > maxBytes) {
    throw new Error(`Remote file is too large. Max size is ${process.env.API_MAX_URL_MB || 25}MB.`);
  }

  return {
    data: buffer.toString('base64'),
    mimeType: response.headers.get('content-type')?.split(';')[0] || 'application/octet-stream',
  };
}

async function inlineInputToPart(input: InlineInput): Promise<GeminiPart> {
  if (input.dataUrl) {
    const parsed = stripDataPrefix(input.dataUrl);
    if (!parsed) {
      throw new Error(`Invalid dataUrl${input.name ? ` for ${input.name}` : ''}.`);
    }
    return {
      inlineData: {
        data: parsed.data,
        mimeType: input.mimeType || parsed.mimeType,
      },
    };
  }

  if (input.url) {
    const fetched = await fetchAsBase64(input.url);
    return {
      inlineData: {
        data: fetched.data,
        mimeType: input.mimeType || fetched.mimeType,
      },
    };
  }

  const rawData = input.base64 || input.data;
  if (!rawData) {
    throw new Error(`File input${input.name ? ` ${input.name}` : ''} is missing data, base64, dataUrl, or url.`);
  }

  const parsed = stripDataPrefix(rawData);
  return {
    inlineData: {
      data: parsed?.data || rawData,
      mimeType: input.mimeType || parsed?.mimeType || 'application/octet-stream',
    },
  };
}

function uploadedFileToPart(file: Express.Multer.File): GeminiPart {
  return {
    inlineData: {
      data: file.buffer.toString('base64'),
      mimeType: file.mimetype || 'application/octet-stream',
    },
  };
}

async function buildParts(body: GenerateRequest, files: Express.Multer.File[] = []) {
  const prompt = (body.prompt || body.instruction || '').trim();
  if (!prompt) {
    throw new Error('prompt or instruction is required.');
  }

  const inlineInputs = [
    ...(body.images || []),
    ...(body.documents || []),
    ...(body.files || []),
  ];

  const parts: GeminiPart[] = [];
  for (const input of inlineInputs) {
    parts.push(await inlineInputToPart(input));
  }

  for (const file of files) {
    parts.push(uploadedFileToPart(file));
  }

  parts.push({ text: prompt });
  return parts;
}

function normalizeMultipartBody(req: express.Request): GenerateRequest {
  return {
    prompt: String(req.body.prompt || req.body.instruction || ''),
    model: req.body.model,
    imageSize: req.body.imageSize,
    aspectRatio: req.body.aspectRatio,
  };
}

function collectResponseParts(response: any) {
  const images: Array<{ mimeType: string; data: string; dataUrl: string }> = [];
  const texts: string[] = [];

  for (const candidate of response.candidates || []) {
    for (const part of candidate.content?.parts || []) {
      if (part.inlineData?.data) {
        const mimeType = part.inlineData.mimeType || 'image/png';
        images.push({
          mimeType,
          data: part.inlineData.data,
          dataUrl: `data:${mimeType};base64,${part.inlineData.data}`,
        });
      }
      if (part.text) {
        texts.push(part.text);
      }
    }
  }

  return { images, texts };
}

async function runGemini(body: GenerateRequest, files: Express.Multer.File[] = [], textOnly = false) {
  const parts = await buildParts(body, files);
  const model = body.model || (textOnly ? defaultTextModel : defaultImageModel);
  const config: Record<string, unknown> = {};

  if (!textOnly) {
    config.imageConfig = {
      imageSize: body.imageSize || '1K',
      aspectRatio: body.aspectRatio || '1:1',
    };
  }

  const response = await ai.models.generateContent({
    model,
    contents: { parts } as any,
    config,
  } as any);

  const { images, texts } = collectResponseParts(response);

  return {
    ok: true,
    model,
    input: {
      fileCount: files.length + (body.images?.length || 0) + (body.documents?.length || 0) + (body.files?.length || 0),
      partCount: parts.length,
      imageSize: textOnly ? undefined : body.imageSize || '1K',
      aspectRatio: textOnly ? undefined : body.aspectRatio || '1:1',
    },
    images,
    texts,
    usageMetadata: response.usageMetadata,
  };
}

function handleError(res: express.Response, error: unknown) {
  const message = error instanceof Error ? error.message : String(error);
  res.status(400).json({
    ok: false,
    error: message,
  });
}

app.get('/api', (_req, res) => {
  res.json({
    ok: true,
    name: 'Gemini Draw API',
    endpoints: [
      'GET /api/health',
      'POST /api/generate',
      'POST /api/generate/upload',
      'POST /api/edit-image',
      'POST /api/analyze-document',
    ],
  });
});

app.get('/api/health', (_req, res) => {
  res.json({
    ok: true,
    apiPort,
    hasGeminiApiKey: Boolean(apiKey),
    authEnabled: Boolean(apiAuthToken),
    defaultImageModel,
    defaultTextModel,
  });
});

app.post('/api/generate', requireAuth, async (req, res) => {
  try {
    res.json(await runGemini(req.body));
  } catch (error) {
    handleError(res, error);
  }
});

app.post('/api/generate/upload', requireAuth, upload.any(), async (req, res) => {
  try {
    res.json(await runGemini(normalizeMultipartBody(req), (req.files as Express.Multer.File[]) || []));
  } catch (error) {
    handleError(res, error);
  }
});

app.post('/api/edit-image', requireAuth, upload.any(), async (req, res) => {
  try {
    const files = (req.files as Express.Multer.File[]) || [];
    const body = req.is('multipart/form-data') ? normalizeMultipartBody(req) : req.body;
    res.json(await runGemini(body, files));
  } catch (error) {
    handleError(res, error);
  }
});

app.post('/api/analyze-document', requireAuth, upload.any(), async (req, res) => {
  try {
    const files = (req.files as Express.Multer.File[]) || [];
    const body = req.is('multipart/form-data') ? normalizeMultipartBody(req) : req.body;
    res.json(await runGemini(body, files, true));
  } catch (error) {
    handleError(res, error);
  }
});

app.listen(apiPort, '0.0.0.0', () => {
  console.log(`Gemini Draw API listening on http://localhost:${apiPort}`);
});