#pragma once
// pipeline-synth-impl.h: private state of the synthesis pipeline
//
// Holds the types that the public API exposes as opaque handles (AceSynth,
// AceSynthJob) plus SynthState, the transient state that flows between the
// primitive ops during a single job.
//
// Included by the two implementation files:
//   pipeline-synth.cpp      orchestrator: load, phases, free
//   pipeline-synth-ops.cpp  primitives:   encode, context, noise, dit, vae

#include "debug.h"
#include "model-store.h"
#include "pipeline-synth.h"
#include "request.h"
#include "timer.h"

#include <string>
#include <vector>

// AceSynth is a thin handle over a ModelStore. It carries the ModelKeys the
// ops use to acquire modules and a pointer to the DiT metadata owned by the
// store (silence_latent, null_cond, config, is_turbo). No GPU module is ever
// owned here: each op performs its own require/release against the store.
struct AceSynth {
    ModelStore *   store;
    AceSynthParams params;

    // CPU metadata pointer. Owned by the store, valid for the store lifetime
    // (always longer than AceSynth). Gives ops access to silence_full,
    // null_cond_cpu, is_turbo and the DiT config without loading the DiT.
    const DiTMeta * meta;

    // Derived constants mirrored for inline use in ops.
    int Oc;      // out_channels (64)
    int ctx_ch;  // in_channels - Oc (128)

    // ModelKeys for the seven GPU modules the pipeline touches.
    ModelKey text_enc_key;   // Qwen3 text encoder, from text_encoder_path
    ModelKey cond_enc_key;   // condition encoder, from dit_path
    ModelKey fsq_tok_key;    // FSQ tokenizer, from dit_path
    ModelKey fsq_detok_key;  // FSQ detokenizer, from dit_path
    ModelKey dit_key;        // DiT, from dit_path (+ adapter_path, adapter_scale)
    ModelKey vae_enc_key;    // VAE encoder, from vae_path
    ModelKey vae_dec_key;    // VAE decoder, from vae_path

    // PP-VAE: optional post-processing VAE (auto-detected from models dir)
    ModelKey pp_vae_enc_key;      // PP-VAE encoder (GGML), from pp_vae_path
    ModelKey pp_vae_dec_key;      // PP-VAE decoder (GGML), from pp_vae_path
    ModelKey pp_vae_enc_ort_key;  // PP-VAE encoder (ORT), from pp_vae_onnx_enc_path
    ModelKey pp_vae_dec_ort_key;  // PP-VAE decoder (ORT), from pp_vae_onnx_dec_path
    bool     have_pp_vae;         // true if PP-VAE GGUF was found
    std::string pp_vae_onnx_enc_path;  // path to scragvae_encoder.onnx (empty = not available)
    std::string pp_vae_onnx_dec_path;  // path to scragvae_decoder.onnx (empty = not available)

    // ORT VAE: optional ONNX Runtime VAE decoder (TensorRT/CUDA EP)
    ModelKey    vae_dec_ort_key;   // VAE-Dec-ORT, from onnx_vae_path
    std::string onnx_vae_path;    // path to vae_decoder.onnx (empty = not available)

    // ORT text/cond encoder: used when DiT is ONNX (.onnx path)
    ModelKey    text_enc_ort_key;   // TextEnc-ORT, from text_encoder.onnx
    ModelKey    cond_enc_ort_key;   // CondEnc-ORT, from cond_encoder.onnx
    bool        is_onnx_pipeline;  // true when DiT is ONNX → all sub-models use ORT
};

// Transient state for a single job, shared by reference across the primitive
// ops that make up phase 1 and phase 2.
struct SynthState {
    // model dimensions (from ctx)
    int Oc;
    int ctx_ch;

    // src audio encoding
    bool               have_cover;
    std::vector<float> cover_latents;        // [T_cover * 64] time-major VAE latents
    std::vector<float> noise_blend_latents;  // clean VAE latents for cover_noise_strength blending
    int                T_cover;

    // outpainting: silence padding for repaint beyond source bounds
    float              left_pad_sec;    // seconds of silence prepended (coordinate shift)
    std::vector<float> padded_src;      // interleaved stereo buffer with silence padding (audio path)
    std::vector<float> padded_latents;  // [T_padded, 64] f32 with silence_full padding (latent path)

    // mode flags
    bool  is_repaint;
    bool  is_lego_region;
    float rs;
    float re;
    bool  use_source_context;
    bool  have_codes;
    int   max_codes_len;

    // audio_codes parsed once per request in ops_resolve_params, reused by
    // ops_build_context. Indexed [0..batch_n). Empty entries mean no codes.
    std::vector<std::vector<int>> per_codes;

    // shared params (from reqs[0])
    AceRequest rr;
    float      duration;
    int        num_steps;
    float      guidance_scale;
    float      shift;

    // diffusion schedule
    std::vector<float> schedule;

    // SDE mode: inject fresh noise at each denoising step (vs ODE pure integration)
    bool use_sde;

    // per-batch seeds (for reproducible SDE re-noising: seed + step offset)
    std::vector<int64_t> seeds;

    // latent dimensions
    int T;
    int S;
    int enc_S;

    // region
    int repaint_t0;
    int repaint_t1;

    // DiT instruction
    std::string instruction_str;  // main DiT instruction (set by orchestrator)

    // conditioning tensors
    std::vector<float> timbre_feats;
    int                S_ref_timbre;  // actual timbre sequence length (from VAE encode)

    // text encoding (per-batch)
    std::vector<std::vector<float>> per_enc;
    std::vector<int>                per_enc_S;
    std::vector<std::vector<float>> per_enc_nc;
    std::vector<int>                per_enc_S_nc;
    bool                            need_enc_switch;

    // stacked encoder hidden states
    int                max_enc_S;
    std::vector<float> enc_hidden;
    std::vector<float> enc_hidden_nc;
    std::vector<int>   per_enc_S_nc_final;
    std::vector<float> null_cond_vec;
    std::vector<float> neg_text_hidden;  // Phase A: text encoder output for negative prompt
    int                neg_S_text = 0;   // token count for negative text

    // DiT context
    std::vector<float> context;
    std::vector<float> context_silence;
    int                cover_steps;

    // noise + output
    std::vector<float> noise;
    std::vector<float> output;
    std::vector<int>   per_S;

    // LRC alignment data (captured during text encoding, consumed by Phase 3)
    std::vector<int>         lyric_token_ids;    // raw tokenized lyric IDs (full lyric prompt)
    std::vector<std::string> lyric_token_texts;  // per-token decoded text (incremental)
    int                      lyric_start_idx;    // pure lyric start (after header tokens)
    int                      lyric_end_idx;      // pure lyric end (before <|endoftext|>)
    std::string              vocal_language;     // "en", "zh", etc. (from request)
    bool                     get_lrc = false;     // whether LRC was requested

    // LRC results (computed inline during ops_dit_generate while DiT is held)
    std::vector<std::string> lrc_results;         // per-batch LRC text (empty if not computed)
    bool                     lrc_done = false;    // true after ops_lrc_extract ran

    // debug / timing
    DebugDumper dbg;
    Timer       timer;
};

// Job bridges phase 1 (DiT) and phase 2 (VAE). The latents live in state.output
// as time-major [batch_n, T, Oc=64] f32 until ace_synth_job_run_vae consumes
// them.
struct AceSynthJob {
    SynthState state;
    int        batch_n;
};