app.py

"""
SonicForge - Flask Main Application

Lightweight local AI music generation workstation.
Entry point for the web application.
"""

import argparse
import logging
import sys
import os

# Ensure the project root is in the Python path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from flask import Flask, render_template, send_from_directory
from flask_cors import CORS

from config import SonicForgeConfig, ModelConfig
from models.base import MusicModel, GenerationParams, GenerationResult
from models.manager import ModelManager
from core.generator import MusicGenerator
from core.exporter import Exporter
from api.routes import api_bp

import numpy as np


# ─── Built-in Demo Models ───────────────────────────────────────────────────


class SineWaveModel(MusicModel):
    """Simple sine wave generator for testing and prototyping."""

    # Note frequencies for musical keys
    KEY_FREQUENCIES = {
        "C": 261.63, "C#": 277.18, "D": 293.66, "D#": 311.13,
        "E": 329.63, "F": 349.23, "F#": 369.99, "G": 392.00,
        "G#": 415.30, "A": 440.00, "A#": 466.16, "B": 493.88,
        "Cm": 261.63, "Dm": 293.66, "Em": 329.63, "Fm": 349.23,
        "Gm": 392.00, "Am": 440.00, "Bm": 493.88,
    }

    GENRE_OSCILLATOR = {
        "Electronic": "sawtooth",
        "Ambient": "sine",
        "Classical": "sine",
        "Pop": "triangle",
        "Rock": "square",
        "Jazz": "sine",
        "Hip-Hop": "square",
        "Lo-Fi": "triangle",
    }

    def __init__(self):
        super().__init__(
            name="sine_basic",
            display_name="Sine Wave Basic",
            description="Simple sine wave generator for testing and prototyping",
        )

    def load(self):
        self._loaded = True

    def unload(self):
        self._loaded = False

    def generate(self, params: GenerationParams) -> GenerationResult:
        sample_rate = params.sample_rate
        n_samples = int(sample_rate * params.duration)
        t = np.linspace(0, params.duration, n_samples, endpoint=False)

        # Base frequency from key
        base_freq = self.KEY_FREQUENCIES.get(params.key, 440.0)

        # Adjust frequency based on BPM (higher BPM = slightly higher pitch perception)
        freq_mod = 1.0 + (params.bpm - 120) / 1200.0
        base_freq *= freq_mod

        # Generate based on genre style
        osc_type = self.GENRE_OSCILLATOR.get(params.genre, "sine")

        if osc_type == "sine":
            signal = np.sin(2 * np.pi * base_freq * t)
        elif osc_type == "square":
            signal = np.sign(np.sin(2 * np.pi * base_freq * t)) * 0.6
        elif osc_type == "sawtooth":
            signal = 2.0 * (base_freq * t - np.floor(0.5 + base_freq * t)) * 0.6
        elif osc_type == "triangle":
            signal = 2.0 * np.abs(2.0 * (base_freq * t - np.floor(base_freq * t + 0.5))) - 1.0
            signal *= 0.7
        else:
            signal = np.sin(2 * np.pi * base_freq * t)

        # Add harmonics for richness
        harmonic_2 = 0.3 * np.sin(2 * np.pi * base_freq * 2 * t)
        harmonic_3 = 0.15 * np.sin(2 * np.pi * base_freq * 3 * t)
        harmonic_5 = 0.08 * np.sin(2 * np.pi * base_freq * 5 * t)
        signal = signal + harmonic_2 + harmonic_3 + harmonic_5

        # Add subtle vibrato
        vibrato = 0.02 * np.sin(2 * np.pi * 5.0 * t)
        signal *= (1.0 + vibrato)

        # Apply amplitude envelope
        attack = 0.05
        release = 0.1
        envelope = np.ones(n_samples)
        attack_samples = int(attack * sample_rate)
        release_samples = int(release * sample_rate)
        if attack_samples > 0:
            envelope[:attack_samples] = np.linspace(0, 1, attack_samples)
        if release_samples > 0 and release_samples < n_samples:
            envelope[-release_samples:] = np.linspace(1, 0, release_samples)
        signal *= envelope

        # Temperature affects harmonic content
        if params.temperature > 0.8:
            noise = np.random.normal(0, 0.02 * (params.temperature - 0.8) * 5, n_samples)
            signal += noise

        # Normalize
        max_val = np.max(np.abs(signal))
        if max_val > 0:
            signal = signal / max_val * 0.85

        return GenerationResult(
            audio_data=signal.astype(np.float64),
            sample_rate=sample_rate,
            duration=params.duration,
            format="wav",
            metadata={
                "model": self._name,
                "prompt": params.prompt,
                "bpm": params.bpm,
                "key": params.key,
                "genre": params.genre,
                "instrument": params.instrument,
                "temperature": params.temperature,
            },
        )


class HarmonicSynthesizerModel(MusicModel):
    """Multi-harmonic synthesizer with chord progression support."""

    CHORD_INTERVALS = {
        "C": [0, 4, 7], "C#": [0, 4, 7], "D": [0, 4, 7], "D#": [0, 4, 7],
        "E": [0, 4, 7], "F": [0, 4, 7], "F#": [0, 4, 7], "G": [0, 4, 7],
        "G#": [0, 4, 7], "A": [0, 4, 7], "A#": [0, 4, 7], "B": [0, 4, 7],
        "Cm": [0, 3, 7], "Dm": [0, 3, 7], "Em": [0, 3, 7], "Fm": [0, 3, 7],
        "Gm": [0, 3, 7], "Am": [0, 3, 7], "Bm": [0, 3, 7],
    }

    BASE_FREQ = 261.63  # Middle C

    def __init__(self):
        super().__init__(
            name="harmonic_synthesizer",
            display_name="Harmonic Synthesizer",
            description="Multi-harmonic synthesizer with chord progression support",
        )

    def load(self):
        self._loaded = True

    def unload(self):
        self._loaded = False

    def generate(self, params: GenerationParams) -> GenerationResult:
        sample_rate = params.sample_rate
        n_samples = int(sample_rate * params.duration)
        t = np.linspace(0, params.duration, n_samples, endpoint=False)

        # Get chord frequencies
        intervals = self.CHORD_INTERVALS.get(params.key, [0, 4, 7])
        chord_freqs = [self.BASE_FREQ * (2 ** (iv / 12.0)) for iv in intervals]

        # BPM determines chord change rate
        beats_per_sec = params.bpm / 60.0
        chord_duration = 4.0 / beats_per_sec  # 4 beats per chord

        signal = np.zeros(n_samples, dtype=np.float64)

        # Generate chord progression (I - IV - V - I)
        progression = [0, 5, 7, 0]  # Semitone offsets for I-IV-V-I
        num_chords = int(np.ceil(params.duration / chord_duration))

        for i in range(num_chords):
            start = int(i * chord_duration * sample_rate)
            end = min(int((i + 1) * chord_duration * sample_rate), n_samples)
            if start >= n_samples:
                break

            chord_t = t[start:end] - t[start]
            semitone_offset = progression[i % len(progression)]

            for freq in chord_freqs:
                shifted_freq = freq * (2 ** (semitone_offset / 12.0))
                chord_signal = np.sin(2 * np.pi * shifted_freq * chord_t)
                # Add warmth with slight detuning
                chord_signal += 0.3 * np.sin(2 * np.pi * shifted_freq * 1.002 * chord_t)
                signal[start:end] += chord_signal

        # Apply envelope per chord section
        for i in range(num_chords):
            start = int(i * chord_duration * sample_rate)
            end = min(int((i + 1) * chord_duration * sample_rate), n_samples)
            if start >= n_samples:
                break

            section_len = end - start
            env = np.ones(section_len)
            attack = min(int(0.02 * sample_rate), section_len // 4)
            release = min(int(0.05 * sample_rate), section_len // 4)
            if attack > 0:
                env[:attack] = np.linspace(0, 1, attack)
            if release > 0:
                env[-release:] = np.linspace(1, 0, release)
            signal[start:end] *= env

        # Normalize
        max_val = np.max(np.abs(signal))
        if max_val > 0:
            signal = signal / max_val * 0.85

        return GenerationResult(
            audio_data=signal.astype(np.float64),
            sample_rate=sample_rate,
            duration=params.duration,
            format="wav",
            metadata={
                "model": self._name,
                "prompt": params.prompt,
                "bpm": params.bpm,
                "key": params.key,
                "genre": params.genre,
                "instrument": params.instrument,
                "temperature": params.temperature,
            },
        )


class RhythmEngineModel(MusicModel):
    """Percussion and rhythm pattern generator."""

    def __init__(self):
        super().__init__(
            name="rhythm_engine",
            display_name="Rhythm Engine",
            description="Percussion and rhythm pattern generator",
        )

    def load(self):
        self._loaded = True

    def unload(self):
        self._loaded = False

    def _generate_kick(self, t, freq_start=150, freq_end=40):
        """Generate a kick drum sound."""
        duration = len(t) / 44100 if len(t) > 0 else 0.1
        freq_sweep = np.linspace(freq_start, freq_end, len(t))
        phase = 2 * np.pi * np.cumsum(freq_sweep) / 44100
        signal = np.sin(phase)
        # Amplitude envelope
        env = np.exp(-t * 30)
        return signal * env * 0.9

    def _generate_snare(self, t):
        """Generate a snare drum sound."""
        noise = np.random.normal(0, 0.5, len(t))
        tone = np.sin(2 * np.pi * 200 * t)
        env = np.exp(-t * 25)
        return (noise * 0.6 + tone * 0.4) * env

    def _generate_hihat(self, t):
        """Generate a hi-hat sound."""
        noise = np.random.normal(0, 0.3, len(t))
        env = np.exp(-t * 60)
        return noise * env

    def generate(self, params: GenerationParams) -> GenerationResult:
        sample_rate = params.sample_rate
        n_samples = int(sample_rate * params.duration)
        signal = np.zeros(n_samples, dtype=np.float64)

        beats_per_sec = params.bpm / 60.0
        beat_samples = int(sample_rate / beats_per_sec)

        # Generate a basic drum pattern
        pattern_length = 4  # 4 beats per bar
        total_beats = int(params.duration * beats_per_sec)

        for beat in range(total_beats):
            beat_pos = beat % pattern_length
            start = int(beat * beat_samples)
            end = min(start + beat_samples, n_samples)
            if start >= n_samples:
                break

            t_beat = np.arange(end - start) / sample_rate

            # Kick on beats 0 and 2
            if beat_pos in (0, 2):
                kick_len = min(int(0.15 * sample_rate), end - start)
                kick = self._generate_kick(t_beat[:kick_len])
                signal[start:start + kick_len] += kick

            # Snare on beats 1 and 3
            if beat_pos in (1, 3):
                snare_len = min(int(0.1 * sample_rate), end - start)
                snare = self._generate_snare(t_beat[:snare_len])
                signal[start:start + snare_len] += snare

            # Hi-hat on every 8th note
            half_beat = beat_samples // 2
            for sub in range(2):
                h_start = start + sub * half_beat
                h_end = min(h_start + int(0.05 * sample_rate), n_samples)
                if h_start < n_samples and h_end > h_start:
                    h_len = h_end - h_start
                    t_hh = np.arange(h_len) / sample_rate
                    hh = self._generate_hihat(t_hh)
                    signal[h_start:h_end] += hh * 0.5

        # Add bass line on kick beats
        bass_freq = 55.0  # Low A
        for beat in range(total_beats):
            if beat % 4 in (0, 2):
                start = int(beat * beat_samples)
                bass_len = min(int(0.3 * sample_rate), n_samples - start)
                if bass_len > 0:
                    t_bass = np.arange(bass_len) / sample_rate
                    bass = 0.4 * np.sin(2 * np.pi * bass_freq * t_bass)
                    env = np.exp(-t_bass * 5)
                    signal[start:start + bass_len] += bass * env

        # Normalize
        max_val = np.max(np.abs(signal))
        if max_val > 0:
            signal = signal / max_val * 0.85

        return GenerationResult(
            audio_data=signal.astype(np.float64),
            sample_rate=sample_rate,
            duration=params.duration,
            format="wav",
            metadata={
                "model": self._name,
                "prompt": params.prompt,
                "bpm": params.bpm,
                "key": params.key,
                "genre": params.genre,
                "instrument": params.instrument,
                "temperature": params.temperature,
            },
        )


# ─── Application Factory ────────────────────────────────────────────────────


def create_app(config: SonicForgeConfig = None) -> Flask:
    """Create and configure the Flask application.

    Args:
        config: Optional SonicForgeConfig instance. If None, loads from environment.

    Returns:
        Configured Flask application instance.
    """
    if config is None:
        config = SonicForgeConfig.from_env()

    app = Flask(
        __name__,
        template_folder=os.path.join(os.path.dirname(__file__), "templates"),
        static_folder=os.path.join(os.path.dirname(__file__), "static"),
    )

    # Enable CORS
    CORS(app)

    # Store config in app
    app.config["APP_CONFIG"] = config
    app.config["SECRET_KEY"] = "sonicforge-secret-key-change-in-production"

    # Initialize model manager
    model_manager = ModelManager(
        auto_unload_minutes=config.auto_unload_minutes,
        max_loaded_models=config.max_loaded_models,
    )

    # Register built-in models
    model_manager.register(SineWaveModel())
    model_manager.register(HarmonicSynthesizerModel())
    model_manager.register(RhythmEngineModel())

    # Load the default model
    model_manager.load_model("sine_basic")

    # Discover plugins
    plugin_dir = os.path.join(os.path.dirname(__file__), "models", "plugins")
    model_manager.add_plugin_dir(plugin_dir)
    model_manager.discover_plugins()

    app.config["MODEL_MANAGER"] = model_manager

    # Initialize music generator
    generator = MusicGenerator(
        model_manager=model_manager,
        output_path=config.output_path,
    )
    app.config["GENERATOR"] = generator

    # Initialize exporter
    exporter = Exporter(output_dir=config.output_path)
    app.config["EXPORTER"] = exporter

    # Register API blueprint
    app.register_blueprint(api_bp)

    # ─── Page Routes ────────────────────────────────────────────────────

    @app.route("/")
    def index():
        """Serve the main application page."""
        return render_template("index.html")

    # ─── Error Handlers ─────────────────────────────────────────────────

    @app.errorhandler(404)
    def not_found(e):
        return {"error": "Not found"}, 404

    @app.errorhandler(500)
    def server_error(e):
        return {"error": "Internal server error"}, 500

    return app


# ─── CLI Entry Point ────────────────────────────────────────────────────────


def main():
    """Main entry point for running the SonicForge server."""
    parser = argparse.ArgumentParser(
        description="SonicForge - Lightweight Local AI Music Generation Workstation",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--host",
        default=None,
        help="Host address to bind to (default: from config or 0.0.0.0)",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=None,
        help="Port to listen on (default: from config or 7860)",
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        default=None,
        help="Enable debug mode",
    )

    args = parser.parse_args()

    # Load configuration
    config = SonicForgeConfig.from_env()

    # Override with CLI arguments
    if args.host is not None:
        config.host = args.host
    if args.port is not None:
        config.port = args.port
    if args.debug is not None:
        config.debug = args.debug

    # Configure logging
    log_level = logging.DEBUG if config.debug else logging.INFO
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    logger = logging.getLogger("SonicForge")
    logger.info("Starting SonicForge v1.0.0")
    logger.info("Host: %s | Port: %d | Debug: %s", config.host, config.port, config.debug)

    # Create and run the application
    app = create_app(config)

    app.run(
        host=config.host,
        port=config.port,
        debug=config.debug,
        use_reloader=False,
    )


if __name__ == "__main__":
    main()