import spaces
import gradio as gr
import torch
import os, json, time, uuid, datetime, pathlib

# torch >= 2.6 defaults torch.load(weights_only=True), which breaks TTS's
# pickle-based config loading. Restore the older behavior before TTS imports.
_orig_torch_load = torch.load
def _patched_torch_load(*args, **kwargs):
    kwargs.setdefault("weights_only", False)
    return _orig_torch_load(*args, **kwargs)
torch.load = _patched_torch_load

os.environ["COQUI_TOS_AGREED"] = "1"

from TTS.api import TTS

device = "cuda" if torch.cuda.is_available() else "cpu"
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)

# ── survey storage: one JSON line per event, batched into a PRIVATE dataset repo ──
# Events: "generate" (Generate clicked; the survey panel opens at the same moment)
#         "submit"   (survey answered). Nothing else is stored — no text, no audio.
SURVEY_REPO = os.environ.get("SURVEY_REPO", "tonyassi/voice-clone-survey")
DATA_DIR = pathlib.Path("survey_data"); DATA_DIR.mkdir(exist_ok=True)
DATA_FILE = DATA_DIR / f"events-{uuid.uuid4().hex[:8]}.jsonl"
scheduler = None
if os.environ.get("TOKEN"):
    from huggingface_hub import CommitScheduler
    scheduler = CommitScheduler(repo_id=SURVEY_REPO, repo_type="dataset",
                                folder_path=DATA_DIR, path_in_repo="data",
                                every=5, token=os.environ["TOKEN"], private=True)

def log_event(event, **fields):
    row = {"ts": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
           "event": event, **fields}
    lock = scheduler.lock if scheduler else __import__("threading").Lock()
    with lock:
        with DATA_FILE.open("a") as f:
            f.write(json.dumps(row) + "\n")

# ── the app itself, unchanged ──
@spaces.GPU
def clone(text, audio):
    tts.tts_to_file(text=text, speaker_wav=audio, language="en", file_path="./output.wav")
    gr.Info("Done. If you have 20 seconds, tell me what you'd want from a better version — the panel on the left.", duration=8)
    return "./output.wav"

def open_survey():
    log_event("generate")
    return gr.Sidebar(open=True)

def submit_survey(use, worth, other, email):
    if not (use or worth or (other or "").strip()):
        return gr.update(value="Pick one thing or write a line first."), gr.update()
    log_event("submit", use=use, worth=worth, other=(other or "").strip()[:500],
              email=(email or "").strip()[:200])
    return gr.update(value="Thank you. That helps more than you'd think."), gr.update(interactive=False)

CSS = """
.top-line { display:block; text-align:center; padding:10px 14px; background:#0f766e; color:#fff;
            text-decoration:none; font-weight:600; cursor:pointer; }
.top-line:hover { filter:brightness(1.05); }
#survey-btn { width:100%; }
/* one choice per line in the survey */
#q-use .wrap, #q-worth .wrap { flex-direction: column; align-items: stretch; }
#q-use label, #q-worth label { width: 100%; }
"""

with gr.Blocks(title="Voice Clone", css=CSS,
               theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate")) as demo:

    top = gr.Button("I'm building a better version of this. 20 seconds — tell me what you'd want →",
                    elem_classes=["top-line"], variant="primary")

    gr.Markdown("""
# Voice Clone
by [Tony Assi](https://www.tonyassi.com/)

---

<h3>If you like voice clone then try <a href="https://huggingface.co/spaces/tonyassi/video-face-swap" target="_blank" rel="noopener noreferrer">Video Face Swap</a></h3>

---

This space uses xtts_v2 model. Non-commercial use only. [Coqui Public Model License](https://huggingface.co/coqui/XTTS-v2/blob/main/LICENSE.txt)

Please ❤️ this Space. [Email me](mailto:tony.assi.media@gmail.com).
""")

    with gr.Row():
        with gr.Column():
            text = gr.Textbox(label="Text")
            audio = gr.Audio(type="filepath", label="Voice reference audio file")
            go = gr.Button("Submit", variant="primary")
        with gr.Column():
            out = gr.Audio(type="filepath", label="output")
            ask = gr.Button("Tell me what you'd want from a better version →", elem_id="survey-btn")

    gr.Examples(
        examples=[["Hey! It's me Dorthy, from the Wizard of Oz. Type in whatever you'd like me to say.", "./audio/Wizard-of-Oz-Dorthy.wav"],
                  ["It's me Vito Corleone, from the Godfather. Type in whatever you'd like me to say.", "./audio/Godfather.wav"],
                  ["Hey, it's me Paris Hilton. Type in whatever you'd like me to say.", "./audio/Paris-Hilton.mp3"],
                  ["Hey, it's me Megan Fox from Transformers. Type in whatever you'd like me to say.", "./audio/Megan-Fox.mp3"],
                  ["Hey there, it's me Jeff Goldblum. Type in whatever you'd like me to say.", "./audio/Jeff-Goldblum.mp3"],
                  ["Hey there, it's me Heath Ledger as the Joker. Type in whatever you'd like me to say.", "./audio/Heath-Ledger.mp3"]],
        inputs=[text, audio],
    )

    with gr.Sidebar(open=False) as side:
        gr.Markdown("### A better version is coming.")
        q_use = gr.Radio(["Fun", "Content for social media", "Narration or audiobook", "Work", "Personal"],
                         label="What are you making this for?", elem_id="q-use")
        q_worth = gr.CheckboxGroup(["Better quality", "Longer text", "Other languages", "Save my voice", "Faster", "Nothing, free is enough"],
                                   label="What would make it worth paying for?", elem_id="q-worth")
        q_other = gr.Textbox(label="Anything else?", placeholder="What would you want it to do?", lines=3)
        q_email = gr.Textbox(label="Email (optional) — I'll tell you when it's ready")
        send = gr.Button("Send", variant="primary")
        thanks = gr.Markdown("")

    # Generate opens the panel first (instant), then runs the job.
    go.click(fn=open_survey, inputs=None, outputs=side, queue=False)
    go.click(fn=clone, inputs=[text, audio], outputs=out)
    top.click(fn=open_survey, inputs=None, outputs=side, queue=False)
    ask.click(fn=open_survey, inputs=None, outputs=side, queue=False)
    send.click(fn=submit_survey, inputs=[q_use, q_worth, q_other, q_email], outputs=[thanks, send])

    demo.queue()

if __name__ == "__main__":
    demo.launch()