from __future__ import annotations
from typing import Iterable

import os
from scipy.io.wavfile import write

import gradio as Gradio
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes

theme = Gradio.themes.Monochrome(
	primary_hue="purple",
	secondary_hue="purple",
	neutral_hue="neutral",
	radius_size=Gradio.themes.sizes.radius_sm,
	font=[Gradio.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
)

class PurpleTheme(Base):
	def __init__(
		self,
		*,
		primary_hue: colors.Color | str = colors.purple,
		secondary_hue: colors.Color | str = colors.purple,
		neutral_hue: colors.Color | str = colors.neutral,
		spacing_size: sizes.Size | str = sizes.spacing_md,
		radius_size: sizes.Size | str = sizes.radius_md,
		font: fonts.Font
		| str
		| Iterable[fonts.Font | str] = (
			fonts.GoogleFont("Inter"),
			"ui-sans-serif",
			"sans-serif",
		),
		font_mono: fonts.Font
		| str
		| Iterable[fonts.Font | str] = (
			fonts.GoogleFont("Space Grotesk"),
			"ui-monospace",
			"monospace",
		),
	):
		super().__init__(
			primary_hue=primary_hue,
			secondary_hue=secondary_hue,
			neutral_hue=neutral_hue,
			spacing_size=spacing_size,
			radius_size=radius_size,
			font=font,
			font_mono=font_mono,
		)
		super().set(
			button_primary_background_fill="linear-gradient(90deg, *primary_300, *secondary_400)",
			button_primary_background_fill_hover="linear-gradient(90deg, *primary_200, *secondary_300)",
			button_primary_text_color="white",
			button_primary_background_fill_dark="linear-gradient(90deg, *primary_600, *secondary_800)",
			block_shadow="*shadow_drop_lg",
			button_shadow="*shadow_drop_lg",
			input_background_fill="zinc",
			input_border_color="*secondary_300",
			input_shadow="*shadow_drop",
			input_shadow_focus="*shadow_drop_lg",
		)

custom_theme = PurpleTheme()

def run_demucs(audio):
	os.makedirs("out", exist_ok=True)
	write('test.wav', audio[0], audio[1])
	result = os.system("python3 -m demucs.separate -n htdemucs_ft -d cpu test.wav -o out")
	print(f"Demucs result: {result}")

	# Check if files exist before returning
	files = ["./out/htdemucs_ft/test/vocals.wav",
		"./out/htdemucs_ft/test/bass.wav",
		"./out/htdemucs_ft/test/drums.wav",
		"./out/htdemucs_ft/test/other.wav"]

	for file in files:
		if not os.path.isfile(file):
			print(f"File not found: {file}")
		else:
			print(f"File exists: {file}")

	return files;

title = "Demucs (finetuned_4s)"
description = "<center>Uses the 'canary bleeding-edge' version of Demucs (v4) that introduces the latest Hybrid Transformer model <br/> Heavily inspired from <a href=\"https://huggingface.co/spaces/Thafx/Demucs_v4_2s_HT\">Thafx's Demucs v4 Space</a>, which is based on <a href=\"https://huggingface.co/spaces/akhaliq/demucs\">akhaliq's PIP Demucs Space</a></center>"
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1911.13254' target='_blank'>Music Source Separation in the Waveform Domain</a> | <a href='https://github.com/facebookresearch/demucs' target='_blank'>Github Repo</a></p>"

Gradio.Interface(
    run_demucs,
    Gradio.Audio(type="numpy", label="Input"),
    [Gradio.Audio(type="filepath", label="Vocals", interactive=False),
     Gradio.Audio(type="filepath", label="Bass", interactive=False),
     Gradio.Audio(type="filepath", label="Drums", interactive=False),
     Gradio.Audio(type="filepath", label="Other", interactive=False)],
    title=title,
    description=description,
    article=article,
    theme=custom_theme,
    analytics_enabled=False,
    css=".generating {visibility: hidden}"
).launch(show_api=True)