import { FastifyInstance, FastifyPluginAsync } from 'fastify';
import { forwardAndSpoofOllamaStreamRequest, forwardChatRequest } from '../proxy/forward';
import { logger } from '../utils/logger';

const ollamaRoutes: FastifyPluginAsync = async (server: FastifyInstance) => {
  server.post('/api/chat', async (request, reply) => {
    try {
      const body = request.body as any;

      if (body?.stream === true) {
        const spoofedStream = await forwardAndSpoofOllamaStreamRequest(body, request.headers.authorization);

        reply.raw.setHeader('Content-Type', 'text/event-stream');
        reply.raw.setHeader('Cache-Control', 'no-cache');
        reply.raw.setHeader('Connection', 'keep-alive');

        return reply.send(spoofedStream);
      }

      const response = await forwardChatRequest(body, request.headers.authorization);
      
      reply.status(200).send(response);
    } catch (error: any) {
      logger.error('Error handling /api/chat:', error.message);
      reply.status(500).send({ error: error.message });
    }
  });
};

export default ollamaRoutes;