import gradio as gr
from pydub import AudioSegment
import math
import os
import zipfile

output_dir = "chunks"
os.makedirs(output_dir, exist_ok=True)

def speed_change(sound, speed=1.0):
    sound_with_altered_frame_rate = sound._spawn(sound.raw_data, overrides={
         "frame_rate": int(sound.frame_rate * speed)
      })
    return sound_with_altered_frame_rate.set_frame_rate(sound.frame_rate)

def get_audio_info(file):
    if file is None:
        return "Bạn chưa upload file!"
    audio = AudioSegment.from_file(file.name)
    file_size_bytes = os.path.getsize(file.name)
    file_size_mb = file_size_bytes / 1024 / 1024
    duration_sec = len(audio) / 1000
    duration_min = int(duration_sec // 60)
    duration_remain_sec = int(duration_sec % 60)
    bitrate_kbps = round((file_size_bytes * 8) / duration_sec / 1000)
    info = (
        f"Thông tin file '{os.path.basename(file.name)}':\n"
        f"- Kích thước: {file_size_mb:.2f} MB\n"
        f"- Độ dài: {duration_min} phút {duration_remain_sec} giây\n"
        f"- Bitrate ước tính: {bitrate_kbps} kbps"
    )
    return info

def process_audio(file, parts, bitrate, speed):
    if file is None:
        return "Bạn chưa upload file!", None
    
    audio = AudioSegment.from_file(file.name)
    base_name = os.path.splitext(os.path.basename(file.name))[0]
    
    audio_to_cut = speed_change(audio, speed) if speed != 1.0 else audio
    
    try:
        num_chunks = int(parts)
    except:
        return "Số phần phải là số nguyên!", None
    if num_chunks < 2 or num_chunks > 10:
        return "Số phần phải là số nguyên từ 2 đến 10.", None
    
    chunk_length_ms = len(audio_to_cut) // num_chunks
    
    paths = []
    for i in range(num_chunks):
        start_ms = i * chunk_length_ms
        end_ms = min((i + 1) * chunk_length_ms, len(audio_to_cut))
        chunk = audio_to_cut[start_ms:end_ms]
        
        if i == num_chunks - 1:
            out_name = f"{base_name}end.mp3"
        else:
            out_name = f"{base_name}{i+1}.mp3"
        out_path = os.path.join(output_dir, out_name)
        
        if bitrate is None:
            chunk.export(out_path, format="mp3")
        else:
            chunk.export(out_path, format="mp3", bitrate=bitrate)
        
        paths.append(out_path)

    # Tạo file zip để tải tất cả
    zip_path = os.path.join(output_dir, f"{base_name}_parts.zip")
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for p in paths:
            zipf.write(p, arcname=os.path.basename(p))
    
    info = get_audio_info(file)
    # Trả về info và file zip
    return info, zip_path

with gr.Blocks() as demo:
    gr.Markdown("## Audio Cutter - Cắt file Âm Thanh thành các file nhỏ")
    
    file_input = gr.File(label="Upload File Âm Thanh")
    out_info = gr.Textbox(label="Thông tin file", interactive=False)
    
    parts_input = gr.Number(value=2, label="Chia thành mấy phần?", precision=0, minimum=2, maximum=10)
    
    bitrate_select = gr.Dropdown(
        choices=[("Không đổi", None), ('64k', '64k'), ('96k', '96k'), ('128k', '128k'), ('192k', '192k'), ('256k', '256k')],
        value=None,
        label="Bit rate của file mới"
    )
    speed_slider = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Tốc độ")
    
    btn_process = gr.Button("Cắt & Tải File")
    out_zip = gr.File(label="Download All (zip)")
    
    file_input.change(get_audio_info, inputs=file_input, outputs=out_info)
    btn_process.click(
        process_audio,
        inputs=[file_input, parts_input, bitrate_select, speed_slider],
        outputs=[out_info, out_zip]
    )

demo.launch(share=True)