import type { Express, Request, Response } from "express";
import { createServer, type Server } from "http";
import { WebSocketServer, WebSocket } from "ws";
import multer from "multer";
import { randomUUID } from "crypto";
import { storage, type AudioFile } from "./storage";
import { getAudioInfo, generateWaveform, analyzeVoice, transformVoice } from "./audioProcessor";
import { transformRequestSchema, batchTransformRequestSchema } from "@shared/schema";
import type { TransformJob, BatchJob } from "@shared/schema";

const jobSubscribers = new Map<string, Set<WebSocket>>();

function broadcastJobUpdate(jobId: string, job: TransformJob) {
  const subscribers = jobSubscribers.get(jobId);
  if (subscribers) {
    const message = JSON.stringify({ type: "job_update", data: job });
    subscribers.forEach((ws) => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(message);
      }
    });
  }
}

const upload = multer({
  storage: multer.memoryStorage(),
  limits: {
    fileSize: 50 * 1024 * 1024,
  },
  fileFilter: (req, file, cb) => {
    const validMimeTypes = ["audio/mpeg", "audio/mp3", "audio/x-mp3", "audio/mpeg3", "audio/x-mpeg-3"];
    const isValidMime = validMimeTypes.includes(file.mimetype);
    const isValidExt = file.originalname.toLowerCase().endsWith(".mp3");
    
    if (isValidMime || isValidExt) {
      cb(null, true);
    } else {
      cb(new Error("Only MP3 files are allowed"));
    }
  },
});

export async function registerRoutes(
  httpServer: Server,
  app: Express
): Promise<Server> {

  const wss = new WebSocketServer({ server: httpServer, path: "/ws" });

  wss.on("connection", (ws: WebSocket) => {
    let subscribedJobId: string | null = null;

    ws.on("message", async (message: Buffer) => {
      try {
        const data = JSON.parse(message.toString());
        
        if (data.type === "subscribe" && data.jobId) {
          const jobId = data.jobId as string;
          
          if (subscribedJobId) {
            const oldSubscribers = jobSubscribers.get(subscribedJobId);
            oldSubscribers?.delete(ws);
          }
          
          subscribedJobId = jobId;
          
          if (!jobSubscribers.has(jobId)) {
            jobSubscribers.set(jobId, new Set());
          }
          jobSubscribers.get(jobId)!.add(ws);
          
          ws.send(JSON.stringify({ type: "subscribed", jobId }));
          
          const currentJob = await storage.getTransformJob(jobId);
          if (currentJob && ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: "job_update", data: currentJob }));
          }
        }
      } catch (error) {
        console.error("WebSocket message error:", error);
      }
    });

    ws.on("close", () => {
      if (subscribedJobId) {
        const subscribers = jobSubscribers.get(subscribedJobId);
        subscribers?.delete(ws);
        if (subscribers?.size === 0) {
          jobSubscribers.delete(subscribedJobId);
        }
      }
    });
  });

  app.post("/api/upload", upload.single("audio"), async (req: Request, res: Response) => {
    try {
      if (!req.file) {
        return res.status(400).json({ error: "No file uploaded" });
      }

      const fileId = randomUUID();
      const buffer = req.file.buffer;

      const [audioInfo, waveformData] = await Promise.all([
        getAudioInfo(buffer),
        generateWaveform(buffer),
      ]);

      const audioFile: AudioFile = {
        id: fileId,
        originalName: req.file.originalname,
        buffer,
        mimeType: req.file.mimetype,
        duration: audioInfo.duration,
        sampleRate: audioInfo.sampleRate,
        channels: audioInfo.channels,
        waveformData,
      };

      await storage.saveAudioFile(audioFile);

      res.json({
        success: true,
        fileId,
        fileName: req.file.originalname,
        duration: audioInfo.duration,
        sampleRate: audioInfo.sampleRate,
        channels: audioInfo.channels,
        waveformData,
      });
    } catch (error) {
      console.error("Upload error:", error);
      res.status(500).json({ error: "Failed to process audio file" });
    }
  });

  app.get("/api/audio/:id", async (req: Request, res: Response) => {
    try {
      const audioFile = await storage.getAudioFile(req.params.id);

      if (!audioFile) {
        return res.status(404).json({ error: "Audio file not found" });
      }

      res.set({
        "Content-Type": audioFile.mimeType,
        "Content-Length": audioFile.buffer.length.toString(),
        "Content-Disposition": `inline; filename="${audioFile.originalName}"`,
      });

      res.send(audioFile.buffer);
    } catch (error) {
      console.error("Audio fetch error:", error);
      res.status(500).json({ error: "Failed to fetch audio file" });
    }
  });

  app.post("/api/transform", async (req: Request, res: Response) => {
    try {
      const validation = transformRequestSchema.safeParse(req.body);

      if (!validation.success) {
        return res.status(400).json({ error: "Invalid request", details: validation.error });
      }

      const { sourceFileId, targetFileId, intensity } = validation.data;

      const sourceFile = await storage.getAudioFile(sourceFileId);
      const targetFile = await storage.getAudioFile(targetFileId);

      if (!sourceFile || !targetFile) {
        return res.status(404).json({ error: "Source or target file not found" });
      }

      const job = await storage.createTransformJob(sourceFileId, targetFileId, intensity);

      processTransformJob(job.id, sourceFile, targetFile, intensity);

      res.json(job);
    } catch (error) {
      console.error("Transform start error:", error);
      res.status(500).json({ error: "Failed to start transformation" });
    }
  });

  app.get("/api/transform/:id", async (req: Request, res: Response) => {
    try {
      const job = await storage.getTransformJob(req.params.id);

      if (!job) {
        return res.status(404).json({ error: "Transform job not found" });
      }

      res.json(job);
    } catch (error) {
      console.error("Transform status error:", error);
      res.status(500).json({ error: "Failed to get transform status" });
    }
  });

  app.get("/api/result/:id", async (req: Request, res: Response) => {
    try {
      const result = await storage.getResultAudio(req.params.id);

      if (!result) {
        return res.status(404).json({ error: "Result not found" });
      }

      res.set({
        "Content-Type": "audio/mpeg",
        "Content-Length": result.buffer.length.toString(),
        "Content-Disposition": `attachment; filename="transformed_${req.params.id}.mp3"`,
      });

      res.send(result.buffer);
    } catch (error) {
      console.error("Result fetch error:", error);
      res.status(500).json({ error: "Failed to fetch result" });
    }
  });

  app.get("/api/waveform/:id", async (req: Request, res: Response) => {
    try {
      const result = await storage.getResultAudio(req.params.id);

      if (!result) {
        return res.status(404).json({ error: "Result not found" });
      }

      res.json({ waveformData: result.waveformData });
    } catch (error) {
      console.error("Waveform fetch error:", error);
      res.status(500).json({ error: "Failed to fetch waveform" });
    }
  });

  app.post("/api/batch-transform", async (req: Request, res: Response) => {
    try {
      const validation = batchTransformRequestSchema.safeParse(req.body);

      if (!validation.success) {
        return res.status(400).json({ error: "Invalid request", details: validation.error });
      }

      const { sourceFileId, targetFileIds, intensity } = validation.data;

      const sourceFile = await storage.getAudioFile(sourceFileId);
      if (!sourceFile) {
        return res.status(404).json({ error: "Source file not found" });
      }

      const targetFiles: AudioFile[] = [];
      for (const targetId of targetFileIds) {
        const targetFile = await storage.getAudioFile(targetId);
        if (!targetFile) {
          return res.status(404).json({ error: `Target file ${targetId} not found` });
        }
        targetFiles.push(targetFile);
      }

      const batchJob = await storage.createBatchJob(sourceFileId, targetFileIds, intensity);

      processBatchJob(batchJob.id, sourceFile, targetFiles, batchJob.jobIds, intensity);

      res.json(batchJob);
    } catch (error) {
      console.error("Batch transform start error:", error);
      res.status(500).json({ error: "Failed to start batch transformation" });
    }
  });

  app.get("/api/batch/:id", async (req: Request, res: Response) => {
    try {
      const batchJob = await storage.getBatchJob(req.params.id);

      if (!batchJob) {
        return res.status(404).json({ error: "Batch job not found" });
      }

      const jobs: TransformJob[] = [];
      for (const jobId of batchJob.jobIds) {
        const job = await storage.getTransformJob(jobId);
        if (job) {
          jobs.push(job);
        }
      }

      res.json({ batchJob, jobs });
    } catch (error) {
      console.error("Batch status error:", error);
      res.status(500).json({ error: "Failed to get batch status" });
    }
  });

  return httpServer;
}

async function processTransformJob(jobId: string, sourceFile: AudioFile, targetFile: AudioFile, intensity: number = 100) {
  try {
    let updatedJob = await storage.updateTransformJob(jobId, {
      status: "analyzing_source",
      progress: 5,
    });
    if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

    const sourceParams = await analyzeVoice(sourceFile.buffer);

    updatedJob = await storage.updateTransformJob(jobId, {
      status: "analyzing_target",
      progress: 15,
      sourceParameters: sourceParams,
    });
    if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

    const targetParams = await analyzeVoice(targetFile.buffer);

    updatedJob = await storage.updateTransformJob(jobId, {
      status: "extracting_features",
      progress: 25,
      targetParameters: targetParams,
    });
    if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

    const { transformedBuffer, resultParams, matchScore } = await transformVoice(
      sourceFile.buffer,
      targetFile.buffer,
      sourceParams,
      targetParams,
      intensity,
      async (progress, stage) => {
        const baseProgress = 25;
        const scaledProgress = baseProgress + (progress * 0.7);
        const stageJob = await storage.updateTransformJob(jobId, {
          status: stage as TransformJob["status"],
          progress: Math.floor(scaledProgress),
        });
        if (stageJob) broadcastJobUpdate(jobId, stageJob);
      }
    );

    const waveformData = await generateWaveform(transformedBuffer);

    await storage.saveResultAudio(jobId, transformedBuffer, waveformData);

    updatedJob = await storage.updateTransformJob(jobId, {
      status: "completed",
      progress: 100,
      resultParameters: resultParams,
      matchScore,
      resultUrl: `/api/result/${jobId}`,
    });
    if (updatedJob) broadcastJobUpdate(jobId, updatedJob);
  } catch (error) {
    console.error("Transform job error:", error);
    const errorJob = await storage.updateTransformJob(jobId, {
      status: "error",
      progress: 0,
      errorMessage: error instanceof Error ? error.message : "Unknown error occurred",
    });
    if (errorJob) broadcastJobUpdate(jobId, errorJob);
  }
}

async function processBatchJob(
  batchJobId: string, 
  sourceFile: AudioFile, 
  targetFiles: AudioFile[], 
  jobIds: string[], 
  intensity: number = 100
) {
  try {
    await storage.updateBatchJob(batchJobId, { status: "processing" });

    const sourceParams = await analyzeVoice(sourceFile.buffer);

    for (let i = 0; i < targetFiles.length; i++) {
      const targetFile = targetFiles[i];
      const jobId = jobIds[i];

      try {
        let updatedJob = await storage.updateTransformJob(jobId, {
          status: "analyzing_source",
          progress: 5,
          sourceParameters: sourceParams,
        });
        if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

        updatedJob = await storage.updateTransformJob(jobId, {
          status: "analyzing_target",
          progress: 15,
        });
        if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

        const targetParams = await analyzeVoice(targetFile.buffer);

        updatedJob = await storage.updateTransformJob(jobId, {
          status: "extracting_features",
          progress: 25,
          targetParameters: targetParams,
        });
        if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

        const { transformedBuffer, resultParams, matchScore } = await transformVoice(
          sourceFile.buffer,
          targetFile.buffer,
          sourceParams,
          targetParams,
          intensity,
          async (progress, stage) => {
            const baseProgress = 25;
            const scaledProgress = baseProgress + (progress * 0.7);
            const stageJob = await storage.updateTransformJob(jobId, {
              status: stage as TransformJob["status"],
              progress: Math.floor(scaledProgress),
            });
            if (stageJob) broadcastJobUpdate(jobId, stageJob);
          }
        );

        const waveformData = await generateWaveform(transformedBuffer);
        await storage.saveResultAudio(jobId, transformedBuffer, waveformData);

        updatedJob = await storage.updateTransformJob(jobId, {
          status: "completed",
          progress: 100,
          resultParameters: resultParams,
          matchScore,
          resultUrl: `/api/result/${jobId}`,
        });
        if (updatedJob) broadcastJobUpdate(jobId, updatedJob);

        const batchJob = await storage.getBatchJob(batchJobId);
        if (batchJob) {
          await storage.updateBatchJob(batchJobId, { 
            completedJobs: batchJob.completedJobs + 1 
          });
        }
      } catch (error) {
        console.error(`Batch job ${jobId} error:`, error);
        const errorJob = await storage.updateTransformJob(jobId, {
          status: "error",
          progress: 0,
          errorMessage: error instanceof Error ? error.message : "Unknown error occurred",
        });
        if (errorJob) broadcastJobUpdate(jobId, errorJob);

        const batchJob = await storage.getBatchJob(batchJobId);
        if (batchJob) {
          await storage.updateBatchJob(batchJobId, { 
            failedJobs: batchJob.failedJobs + 1 
          });
        }
      }
    }

    const finalBatchJob = await storage.getBatchJob(batchJobId);
    if (finalBatchJob) {
      const allCompleted = finalBatchJob.completedJobs + finalBatchJob.failedJobs === finalBatchJob.totalJobs;
      if (allCompleted) {
        await storage.updateBatchJob(batchJobId, { 
          status: finalBatchJob.failedJobs > 0 ? "error" : "completed" 
        });
      }
    }
  } catch (error) {
    console.error("Batch processing error:", error);
    await storage.updateBatchJob(batchJobId, { status: "error" });
  }
}
