import gradio as gr
import edge_tts
import asyncio
import tempfile
import os
import time
import base64
from datetime import datetime
from pydub import AudioSegment

VOICE      = "en-US-AndrewMultilingualNeural"
CHUNK_SIZE = 5000
HF_USER    = "Cracked-User"
HF_REPO    = "tts-audiobooks"

# ── Exact same logic as Colab ─────────────────────────────────────────────────
def run_tts(text, custom_name=''):
    chunks     = [text[i:i+CHUNK_SIZE] for i in range(0, len(text), CHUNK_SIZE)]
    total      = len(chunks)
    parts      = []
    chunk_times = []
    batch_start = time.time()
    tmp_dir    = tempfile.mkdtemp()

    yield None, '⏳ Total chunks: ' + str(total), make_player(None), ''

    for i, chunk in enumerate(chunks):
        path = tmp_dir + '/chunk_' + str(i).zfill(5) + '.mp3'

        if chunk_times:
            avg     = sum(chunk_times) / len(chunk_times)
            eta     = int(avg * (total - i))
            elapsed = int(time.time() - batch_start)
            status  = (
                'Chunk ' + str(i+1) + '/' + str(total) +
                '  ETA: ' + str(eta//60) + 'm ' + str(eta%60) + 's' +
                '  Elapsed: ' + str(elapsed//60) + 'm ' + str(elapsed%60) + 's'
            )
        else:
            status = 'Chunk ' + str(i+1) + '/' + str(total) + '  ETA: calculating...'

        yield None, status, make_player(None), ''

        try:
            async def save(chunk=chunk, path=path):
                comm = edge_tts.Communicate(chunk, VOICE)
                await comm.save(path)
            asyncio.run(save())
            parts.append(path)
        except Exception as e:
            try:
                time.sleep(5)
                asyncio.run(save())
                parts.append(path)
            except Exception as e2:
                yield None, '❌ Failed chunk ' + str(i+1) + ': ' + str(e2), make_player(None), ''
                return

        chunk_times.append(time.time() - batch_start - sum(chunk_times))

        # Log progress to HF so you can check after refresh
        avg2    = sum(chunk_times) / len(chunk_times)
        eta2    = int(avg2 * (total - i - 1))
        elapsed2 = int(time.time() - batch_start)
        pct2    = int(((i+1) / total) * 100)
        log_msg = (
            '🔊 Chunk ' + str(i+1) + ' / ' + str(total) + '  (' + str(pct2) + '%)\n' +
            '⏱️ Elapsed: ' + str(elapsed2//60) + 'm ' + str(elapsed2%60) + 's\n' +
            '⏳ ETA: ' + str(eta2//60) + 'm ' + str(eta2%60) + 's\n' +
            '🕐 Updated: ' + datetime.now().strftime('%H:%M:%S')
        )
        import threading
        threading.Thread(target=log_progress, args=(log_msg, os.environ.get('HF_TOKEN', '')), daemon=True).start()

    yield None, 'Merging all chunks...', make_player(None), ''

    combined = AudioSegment.empty()
    for p in sorted(parts):
        if os.path.exists(p):
            combined += AudioSegment.from_file(p) + AudioSegment.silent(200)

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    filename  = (custom_name.strip().replace(' ', '_') + '.mp3') if custom_name and custom_name.strip() else ('andrew_' + timestamp + '.mp3')
    out_path  = tempfile.mktemp(suffix='.mp3')
    combined.export(out_path, format='mp3', bitrate='192k')

    # Cleanup chunks
    for p in parts:
        try: os.remove(p)
        except: pass

    total_time = int(time.time() - batch_start)
    done_msg   = 'DONE! Time: ' + str(total_time//60) + 'm ' + str(total_time%60) + 's'

    log_progress('✅ DONE! Uploading audio file...\n🕐 ' + datetime.now().strftime('%H:%M:%S'), os.environ.get('HF_TOKEN', ''))
    yield None, done_msg + '\nUploading to HF...', make_player(None), ''
    hf_msg = upload_to_hf(out_path, filename)

    yield out_path, done_msg, make_player(out_path), hf_msg

# ── Upload to HF Dataset ──────────────────────────────────────────────────────
def upload_to_hf(file_path, filename):
    try:
        from huggingface_hub import HfApi
        token = os.environ.get('HF_TOKEN', '')
        if not token:
            return '⚠️ No HF_TOKEN secret found.'
        api = HfApi()
        # Create repo if it doesn't exist
        try:
            api.create_repo(
                repo_id=HF_USER + '/' + HF_REPO,
                repo_type='dataset',
                private=True,
                token=token,
                exist_ok=True
            )
        except:
            pass
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo='audiobooks/' + filename,
            repo_id=HF_USER + '/' + HF_REPO,
            repo_type='dataset',
            token=token
        )
        url = 'https://huggingface.co/datasets/' + HF_USER + '/' + HF_REPO
        return '✅ Saved!\n🔗 ' + url
    except Exception as e:
        return '❌ Upload error: ' + str(e)

# ── Log progress to HF Dataset ───────────────────────────────────────────────
def log_progress(msg, token):
    try:
        from huggingface_hub import HfApi
        if not token:
            return
        api = HfApi()
        log_path = '/tmp/progress.txt'
        with open(log_path, 'w') as f:
            f.write(msg)
        api.upload_file(
            path_or_fileobj=log_path,
            path_in_repo='progress.txt',
            repo_id=HF_USER + '/' + HF_REPO,
            repo_type='dataset',
            token=token
        )
    except:
        pass

# ── Custom HTML player ────────────────────────────────────────────────────────
def make_player(file_path):
    if not file_path or not os.path.exists(file_path):
        return '<div style="color:#3a3a5a;padding:12px;text-align:center;border:1px dashed #2a2a3a;border-radius:10px;">🎧 Player appears here after generation</div>'
    with open(file_path, 'rb') as f:
        b64 = base64.b64encode(f.read()).decode()
    return '''
<div style="background:#0e0e1c;border:1px solid #2a2a3a;border-radius:14px;padding:20px;margin-top:10px;">
  <div style="color:#ffd700;font-family:sans-serif;font-weight:700;font-size:0.78rem;
              letter-spacing:0.06em;text-transform:uppercase;margin-bottom:12px;">
    🎧 Preview Player
  </div>
  <audio id="tts-player" controls style="width:100%;margin-bottom:14px;accent-color:#ffd700;">
    <source src="data:audio/mpeg;base64,''' + b64 + '''" type="audio/mpeg">
  </audio>
  <div style="display:flex;align-items:center;gap:10px;flex-wrap:wrap;">
    <span style="color:#9ca3af;font-size:0.82rem;">Speed:</span>
    <button onclick="setSpd(this,0.75)" style="background:#1a1a2e;border:1px solid #3a3a5c;border-radius:8px;color:#9ca3af;padding:5px 12px;cursor:pointer;font-size:0.82rem;">0.75x</button>
    <button onclick="setSpd(this,1.0)" style="background:#b8860b;border:none;border-radius:8px;color:white;padding:5px 12px;cursor:pointer;font-size:0.82rem;">1x</button>
    <button onclick="setSpd(this,1.25)" style="background:#1a1a2e;border:1px solid #3a3a5c;border-radius:8px;color:#9ca3af;padding:5px 12px;cursor:pointer;font-size:0.82rem;">1.25x</button>
    <button onclick="setSpd(this,1.5)" style="background:#1a1a2e;border:1px solid #3a3a5c;border-radius:8px;color:#9ca3af;padding:5px 12px;cursor:pointer;font-size:0.82rem;">1.5x</button>
    <button onclick="setSpd(this,1.75)" style="background:#1a1a2e;border:1px solid #3a3a5c;border-radius:8px;color:#9ca3af;padding:5px 12px;cursor:pointer;font-size:0.82rem;">1.75x</button>
    <button onclick="setSpd(this,2.0)" style="background:#1a1a2e;border:1px solid #3a3a5c;border-radius:8px;color:#9ca3af;padding:5px 12px;cursor:pointer;font-size:0.82rem;">2x</button>
  </div>
  <script>
    function setSpd(btn, s) {
      var p = document.getElementById('tts-player');
      if(p) p.playbackRate = s;
      btn.closest('div').querySelectorAll('button').forEach(function(b){
        b.style.background='#1a1a2e'; b.style.border='1px solid #3a3a5c'; b.style.color='#9ca3af';
      });
      btn.style.background='#b8860b'; btn.style.border='none'; btn.style.color='white';
    }
  </script>
</div>'''

def generate_from_text(text, custom_name=''):
    if not text.strip():
        yield None, '❌ Please enter some text!', make_player(None), ''
        return
    yield from run_tts(text)

def generate_from_file(file, custom_name=''):
    if file is None:
        yield None, '❌ Please upload a .txt file!', make_player(None), ''
        return
    try:
        with open(file, 'r', encoding='utf-8') as f:
            text = f.read()
        if not text.strip():
            yield None, '❌ File is empty!', make_player(None), ''
            return
        yield None, '📄 Loaded ' + str(len(text)) + ' chars. Starting...', make_player(None), ''
        yield from run_tts(text, custom_name)
    except Exception as e:
        yield None, '❌ Could not read file: ' + str(e), make_player(None), ''

# ── CSS ───────────────────────────────────────────────────────────────────────
css = """
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@700;800&family=DM+Sans:wght@300;400;500&display=swap');
* { font-family: 'DM Sans', sans-serif; }
body, .gradio-container { background: #07070f !important; color: #e8e8f5 !important; }
.gradio-container { max-width: 900px !important; margin: 0 auto !important; }
h1 { font-family: 'Syne', sans-serif !important; font-weight: 800 !important;
     font-size: 2.2rem !important;
     background: linear-gradient(135deg, #ffd700, #ff8c00, #ff4500) !important;
     -webkit-background-clip: text !important; -webkit-text-fill-color: transparent !important;
     text-align: center !important; margin-bottom: 0.2rem !important; }
.subtitle { text-align:center; color:#5a5a7a; font-size:0.88rem; margin-bottom:1.5rem; }
label { font-family:'Syne',sans-serif !important; font-weight:700 !important;
        font-size:0.78rem !important; color:#ffd700 !important;
        letter-spacing:0.06em !important; text-transform:uppercase !important; }
textarea { background:#0e0e1c !important; border:1px solid #2a2a3a !important;
    border-radius:10px !important; color:#e8e8f5 !important; padding:12px !important; font-size:0.95rem !important; }
textarea:focus { border-color:#ffd700 !important; outline:none !important; }
.block { background:#0d0d1a !important; border:1px solid #1a1a2e !important;
         border-radius:14px !important; padding:1.2rem !important; }
.btn-gen { background:linear-gradient(135deg,#b8860b,#ff8c00) !important;
    border:none !important; border-radius:12px !important; color:white !important;
    font-family:'Syne',sans-serif !important; font-weight:700 !important;
    font-size:1.05rem !important; padding:16px !important; width:100% !important; }
.btn-gen:hover { opacity:0.85 !important; }
.tip { background:#0e0e1c; border-left:3px solid #ffd700; border-radius:0 8px 8px 0;
       padding:10px 14px; font-size:0.82rem; color:#9ca3af; margin:0.5rem 0; }
"""

# ── GUI ───────────────────────────────────────────────────────────────────────
with gr.Blocks(css=css, title='Andrew TTS') as demo:
    gr.HTML('<h1>🎙️ Andrew TTS</h1>')
    gr.HTML('<p class="subtitle">Microsoft Andrew voice · Exact Colab quality · Auto saves to HF Dataset</p>')

    with gr.Tabs():
        with gr.Tab('📝  Paste Text'):
            with gr.Row():
                with gr.Column(scale=2):
                    text_input = gr.Textbox(
                        label='Your Text',
                        placeholder='Paste your novel chapter or any text here — no limit!',
                        lines=12
                    )
                    name1 = gr.Textbox(label='File Name (optional)', placeholder='e.g. chapter_102 — leave blank for auto name')
                    btn1 = gr.Button('🎙️ Generate Audio', elem_classes='btn-gen')
                with gr.Column(scale=1):
                    gr.HTML('<div class="tip">• No character limit<br>• Live progress + ETA<br>• Auto saves to HF Dataset forever<br>• Speed control 0.75x–2x</div>')
                    status1 = gr.Textbox(label='Progress', interactive=False, lines=3)
                    hf1     = gr.Textbox(label='Storage Status', interactive=False, lines=2)
            player1 = gr.HTML(make_player(None))
            audio1  = gr.Audio(label='Download Audio', type='filepath')
            btn1.click(fn=generate_from_text, inputs=[text_input, name1], outputs=[audio1, status1, player1, hf1])

        with gr.Tab('📁  Upload .txt File'):
            with gr.Row():
                with gr.Column(scale=2):
                    file_input = gr.File(label='Upload .txt file', file_types=['.txt'])
                    name2 = gr.Textbox(label='File Name (optional)', placeholder='e.g. chapter_102 — leave blank for auto name')
                    btn2       = gr.Button('🎙️ Generate Audio', elem_classes='btn-gen')
                with gr.Column(scale=1):
                    gr.HTML('<div class="tip">• Upload any .txt file<br>• Live progress + ETA<br>• Auto saves to HF Dataset forever<br>• Speed control 0.75x–2x</div>')
                    status2 = gr.Textbox(label='Progress', interactive=False, lines=3)
                    hf2     = gr.Textbox(label='Storage Status', interactive=False, lines=2)
            player2 = gr.HTML(make_player(None))
            audio2  = gr.Audio(label='Download Audio', type='filepath')
            btn2.click(fn=generate_from_file, inputs=[file_input, name2], outputs=[audio2, status2, player2, hf2])

demo.launch()