import spaces
import zero
import gradio as gr
import sys
import threading
import queue
import time
import random
from io import TextIOBase
import datetime
import subprocess
import os
from inference import postprocess_inst_names
from inference import inference_patch
from convert import abc2xml, xml2, pdf2img


title_html = """
<div class="title-container">
    <h1 class="title-text">NotaGen</h1> &nbsp;
        <!-- ArXiv -->
        <a href="https://arxiv.org/abs/2502.18008">
            <img src="https://img.shields.io/badge/NotaGen_Paper-ArXiv-%23B31B1B?logo=arxiv&logoColor=white" alt="Paper">
        </a>
        &nbsp;
        <!-- GitHub -->
        <a href="https://github.com/ElectricAlexis/NotaGen">
            <img src="https://img.shields.io/badge/NotaGen_Code-GitHub-%23181717?logo=github&logoColor=white" alt="GitHub">
        </a>
        &nbsp;
        <!-- HuggingFace -->
        <a href="https://huggingface.co/ElectricAlexis/NotaGen">
            <img src="https://img.shields.io/badge/NotaGen_Weights-HuggingFace-%23FFD21F?logo=huggingface&logoColor=white" alt="Weights">
        </a>
        &nbsp;
        <!-- Web Demo -->
        <a href="https://electricalexis.github.io/notagen-demo/">
            <img src="https://img.shields.io/badge/NotaGen_Demo-Web-%23007ACC?logo=google-chrome&logoColor=white" alt="Demo">
        </a>
</div>
<bp>
<p style="font-size: 1.2em;">NotaGen is a model for generating sheet music in ABC notation format. Select a period, composer, and instrumentation to generate classical-style music!</p>
"""

# Read prompt combinations
with open('prompts.txt', 'r') as f:
    prompts = f.readlines()

valid_combinations = set()
for prompt in prompts:
    prompt = prompt.strip()
    parts = prompt.split('_')
    valid_combinations.add((parts[0], parts[1], parts[2]))

# Prepare dropdown options
periods = sorted({p for p, _, _ in valid_combinations})
composers = sorted({c for _, c, _ in valid_combinations})
instruments = sorted({i for _, _, i in valid_combinations})

# Dynamically update composer and instrument dropdown options
def update_components(period, composer):
    if not period:
        return [
            gr.update(choices=[], value=None, interactive=False),
            gr.update(choices=[], value=None, interactive=False)
        ]

    valid_composers = sorted({c for p, c, _ in valid_combinations if p == period})
    valid_instruments = sorted({i for p, c, i in valid_combinations if p == period and c == composer}) if composer else []

    return [
        gr.update(
            choices=valid_composers,
            value=composer if composer in valid_composers else None,
            interactive=True
        ),
        gr.update(
            choices=valid_instruments,
            value=None,
            interactive=bool(valid_instruments)
        )
    ]

# Custom realtime stream for outputting model inference process to frontend
class RealtimeStream(TextIOBase):
    def __init__(self, queue):
        self.queue = queue

    def write(self, text):
        self.queue.put(text)
        return len(text)

def convert_files(abc_content, period, composer, instrumentation):
    if not all([period, composer, instrumentation]):
        raise gr.Error("Please complete a valid generation first before saving")

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    prompt_str = f"{period}_{composer}_{instrumentation}"
    filename_base = f"{timestamp}_{prompt_str}"

    abc_filename = f"{filename_base}.abc"
    with open(abc_filename, "w", encoding="utf-8") as f:
        f.write(abc_content)

    # instrumentation replacement
    postprocessed_inst_abc = postprocess_inst_names(abc_content)
    filename_base_postinst = f"{filename_base}_postinst"
    with open(filename_base_postinst + ".abc", "w", encoding="utf-8") as f:
        f.write(postprocessed_inst_abc)

    # Convert files
    file_paths = {'abc': abc_filename}
    try:
        # abc2xml
        abc2xml(filename_base)
        abc2xml(filename_base_postinst)

        # xml2pdf
        xml2(filename_base, 'pdf')

        # xml2mid
        xml2(filename_base, 'mid')
        xml2(filename_base_postinst, 'mid')

        # xml2mp3
        xml2(filename_base, 'mp3')
        xml2(filename_base_postinst, 'mp3')

        # 将PDF转为图片
        images = pdf2img(filename_base)
        for i, image in enumerate(images):
            image.save(f"{filename_base}_page_{i+1}.png", "PNG")

        file_paths.update({
            'xml': f"{filename_base_postinst}.xml",
            'pdf': f"{filename_base}.pdf",
            'mid': f"{filename_base_postinst}.mid",
            'mp3': f"{filename_base_postinst}.mp3",
            'pages': len(images),
            'current_page': 0,
            'base': filename_base
        })

    except Exception as e:
        raise gr.Error(f"File processing failed: {str(e)}")

    return file_paths


# Page navigation control function
def update_page(direction, data):
    """
    data contains three key pieces of information: 'pages', 'current_page', and 'base'
    """
    if not data:
        return None, gr.update(interactive=False), gr.update(interactive=False), data

    if direction == "prev" and data['current_page'] > 0:
        data['current_page'] -= 1
    elif direction == "next" and data['current_page'] < data['pages'] - 1:
        data['current_page'] += 1

    current_page_index = data['current_page']
    # Update image path
    new_image = f"{data['base']}_page_{current_page_index+1}.png"
    # When current_page==0, prev_btn is disabled; when current_page==pages-1, next_btn is disabled
    prev_btn_state = gr.update(interactive=(current_page_index > 0))
    next_btn_state = gr.update(interactive=(current_page_index < data['pages'] - 1))

    return new_image, prev_btn_state, next_btn_state, data


@spaces.GPU#(duration=600)
def generate_music(period, composer, instrumentation):
    """
    Must ensure each yield returns the same number of values.
    We're preparing to return 5 values, corresponding to:
    1) process_output (intermediate inference information)
    2) final_output (final ABC)
    3) pdf_image (path to the PNG of the first page of the PDF)
    4) audio_player (mp3 path)
    5) pdf_state (state for page navigation)
    """
    # Set a different random seed each time based on current timestamp
    random_seed = int(time.time()) % 10000
    random.seed(random_seed)
    
    # For numpy if you're using it
    try:
        import numpy as np
        np.random.seed(random_seed)
    except ImportError:
        pass
        
    # For torch if you're using it
    try:
        import torch
        torch.manual_seed(random_seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(random_seed)
    except ImportError:
        pass

    if (period, composer, instrumentation) not in valid_combinations:
        # If the combination is invalid, raise an error
        raise gr.Error("Invalid prompt combination! Please re-select from the period options")

    output_queue = queue.Queue()
    original_stdout = sys.stdout
    sys.stdout = RealtimeStream(output_queue)

    result_container = []

    def run_inference():
        try:
            # Use downloaded model weights path for inference
            result = inference_patch(period, composer, instrumentation)
            result_container.append(result)
        finally:
            sys.stdout = original_stdout

    thread = threading.Thread(target=run_inference)
    thread.start()

    process_output = ""
    final_output_abc = ""
    pdf_image = None
    audio_file = None
    pdf_state = None

    # First continuously read intermediate output
    while thread.is_alive():
        try:
            text = output_queue.get(timeout=0.1)
            process_output += text
            # No final ABC yet, files not yet converted
            yield process_output, final_output_abc, pdf_image, audio_file, pdf_state, gr.update(value=None, visible=False)
        except queue.Empty:
            continue

    # After thread ends, get all remaining items from the queue
    while not output_queue.empty():
        text = output_queue.get()
        process_output += text

    # Final inference result
    final_result = result_container[0] if result_container else ""
    
    # Display file conversion prompt
    final_output_abc = "Converting files..."
    yield process_output, final_output_abc, pdf_image, audio_file, pdf_state, gr.update(value=None, visible=False)


    # Convert files
    try:
        file_paths = convert_files(final_result, period, composer, instrumentation)
        final_output_abc = final_result
        # Get the first image and mp3 file
        if file_paths['pages'] > 0:
            pdf_image = f"{file_paths['base']}_page_1.png"
        audio_file = file_paths['mp3']
        pdf_state = file_paths  # Directly use the converted information dictionary as state
        
        # Prepare download file list
        download_list = []
        if 'abc' in file_paths and os.path.exists(file_paths['abc']):
            download_list.append(file_paths['abc'])
        if 'xml' in file_paths and os.path.exists(file_paths['xml']):
            download_list.append(file_paths['xml'])
        if 'pdf' in file_paths and os.path.exists(file_paths['pdf']):
            download_list.append(file_paths['pdf'])
        if 'mid' in file_paths and os.path.exists(file_paths['mid']):
            download_list.append(file_paths['mid'])
        if 'mp3' in file_paths and os.path.exists(file_paths['mp3']):
            download_list.append(file_paths['mp3'])
    except Exception as e:
        # If conversion fails, return error message to output box
        yield process_output, f"Error converting files: {str(e)}", None, None, None, gr.update(value=None, visible=False)
        return

   # Final yield with all information - modify here to make component visible
    yield process_output, final_output_abc, pdf_image, audio_file, pdf_state, gr.update(value=download_list, visible=True)


def get_file(file_type, period, composer, instrumentation):
    """
    Returns the local file of specified type for Gradio download
    """
    # Here you actually need to return based on specific file paths saved earlier, simplified for demo
    # If matching by timestamp, you can store all converted files in a directory and get the latest
    # This is just an example:
    possible_files = [f for f in os.listdir('.') if f.endswith(f'.{file_type}')]
    if not possible_files:
        return None
    # Simply return the latest
    possible_files.sort(key=os.path.getmtime)
    return possible_files[-1]


css = """
/* Compact button style */
button[size="sm"] {
    padding: 4px 8px !important;
    margin: 2px !important;
    min-width: 60px;
}

/* PDF preview area */
#pdf-preview {
    border-radius: 8px;  /* Rounded corners */
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);  /* Shadow */
}

.page-btn {
    padding: 12px !important;  /* Increase clickable area */
    margin: auto !important;   /* Vertical center */
}

/* Button hover effect */
.page-btn:hover {
    background: #f0f0f0 !important;
    transform: scale(1.05);
}

/* Layout adjustment */
.gr-row {
    gap: 10px !important;  /* Element spacing */
}

/* Audio player */
.audio-panel {
    margin-top: 15px !important;
    max-width: 400px;
}

#audio-preview audio {
    height: 200px !important;
}

/* Save functionality area */
.save-as-row {
    margin-top: 15px;
    padding: 10px;
    border-top: 1px solid #eee;
}

/* Download files styling */
.download-files {
    margin-top: 15px;
    border-radius: 8px;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

/* Social icons styling */
.title-container {
    display: flex;
    align-items: center;
    gap: 15px;
    margin-bottom: 10px;
}

.title-text {
    margin: 0;
    font-size: 1.8em;
}

.social-icons {
    display: flex;
    gap: 10px;
}

.social-icon {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    width: 32px;
    height: 32px;
    border-radius: 50%;
    background-color: #f5f5f5;
    text-decoration: none;
    transition: transform 0.2s, background-color 0.2s;
}

.social-icon:hover {
    transform: scale(1.1);
    background-color: #e0e0e0;
}

.social-icon img {
    width: 20px;
    height: 20px;
}

"""

with gr.Blocks(css=css) as demo:
    gr.HTML(title_html)

    # For storing PDF page count, current page and other information
    pdf_state = gr.State()

    with gr.Column():
        with gr.Row():
            # Left sidebar
            with gr.Column():
                with gr.Row():
                    period_dd = gr.Dropdown(
                        choices=periods,
                        value=None,
                        label="Period",
                        interactive=True
                    )
                    composer_dd = gr.Dropdown(
                        choices=[],
                        value=None,
                        label="Composer",
                        interactive=False
                    )
                    instrument_dd = gr.Dropdown(
                        choices=[],
                        value=None,
                        label="Instrumentation",
                        interactive=False
                    )

                generate_btn = gr.Button("Generate!", variant="primary")

                process_output = gr.Textbox(
                    label="Generation process",
                    interactive=False,
                    lines=2,
                    max_lines=2,
                    placeholder="Generation progress will be shown here..."
                )

                final_output = gr.Textbox(
                    label="Post-processed ABC notation scores",
                    interactive=True,
                    lines=8,
                    max_lines=8,
                    placeholder="Post-processed ABC scores will be shown here..."
                )

                # Audio playback
                audio_player = gr.Audio(
                    label="Audio Preview",
                    format="mp3",
                    interactive=False,
                )

            # Right sidebar
            with gr.Column():
                # Image container
                pdf_image = gr.Image(
                    label="Sheet Music Preview",
                    show_label=False,
                    height=650,
                    type="filepath",
                    elem_id="pdf-preview",
                    interactive=False,
                    show_download_button=False
                )

                # Page navigation buttons
                with gr.Row():
                    prev_btn = gr.Button(
                        "⬅️ Last Page",
                        variant="secondary",
                        size="sm",
                        elem_classes="page-btn"
                    )
                    next_btn = gr.Button(
                        "Next Page ➡️",
                        variant="secondary",
                        size="sm",
                        elem_classes="page-btn"
                    )

        with gr.Column():
            gr.Markdown("**Download Files:**")
            download_files = gr.Files(
                label="Generated Files", 
                visible=False,
                elem_classes="download-files",
                type="filepath"  # Make sure this is set to filepath
            )

    # Dropdown linking
    period_dd.change(
        update_components,
        inputs=[period_dd, composer_dd],
        outputs=[composer_dd, instrument_dd]
    )
    composer_dd.change(
        update_components,
        inputs=[period_dd, composer_dd],
        outputs=[composer_dd, instrument_dd]
    )

    # Click generate button, note outputs must match each yield in generate_music
    generate_btn.click(
        generate_music,
        inputs=[period_dd, composer_dd, instrument_dd],
        outputs=[process_output, final_output, pdf_image, audio_player, pdf_state, download_files]
    )

    # Page navigation
    prev_signal = gr.Textbox(value="prev", visible=False)
    next_signal = gr.Textbox(value="next", visible=False)

    prev_btn.click(
        update_page,
        inputs=[prev_signal, pdf_state],  # ✅ Use component
        outputs=[pdf_image, prev_btn, next_btn, pdf_state]
    )

    next_btn.click(
        update_page,
        inputs=[next_signal, pdf_state],  # ✅ Use component
        outputs=[pdf_image, prev_btn, next_btn, pdf_state]
    )


if __name__ == "__main__":
    # Configure GPU/CPU handling
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        # share=True
    )