from flask import Flask, request, render_template, jsonify
import os, re
from whisper_tts import WhisperTTS
from ollama_chatbotTTS import OllamaChat
from text_to_speech import TextToSpeech
from sync_audio_video import AudioVideoSync

app = Flask(__name__, static_folder='.', static_url_path='')

THUMBNAILS_DIR = "thumbnails"
VIDEO_DIR = "sample_video"
UPLOAD_DIR = "uploads"

def get_thumbnail_images():
    if not os.path.exists(THUMBNAILS_DIR):
        return []
    return [os.path.splitext(f)[0] for f in os.listdir(THUMBNAILS_DIR) if f.lower().endswith((".png",".jpg",".jpeg"))]

@app.route('/')
def index():
    avatars = get_thumbnail_images()
    return render_template('index.html', avatars=avatars)

@app.route('/transcribe', methods=['POST'])
def transcribe():
    f = request.files.get('audio')
    if f:
        os.makedirs(UPLOAD_DIR, exist_ok=True)
        path = os.path.join(UPLOAD_DIR, f.filename)
        f.save(path)
        text = WhisperTTS().transcribe_audio(path)
        return jsonify(text=text)
    return jsonify(text="")

@app.route('/chat', methods=['POST'])
def chat():
    data = request.get_json()
    text = data.get('text', "")
    r = OllamaChat().get_response(text)
    resp = re.sub(r"<think>|</think>", "", r).strip()
    return jsonify(response=resp)

@app.route('/tts', methods=['POST'])
def tts():
    data = request.get_json()
    text = data.get('text', "")
    if not text:
        return jsonify(audio_url="")
    path = TextToSpeech().synthesize(text)
    return jsonify(audio_url="/" + path)

@app.route('/sync', methods=['POST'])
def sync_audio_video():
    data = request.get_json()
    avatar = data.get('avatar', "")
    audio_url = data.get('audio_url', "")
    audio_path = audio_url.lstrip('/')
    vid = None
    for v in os.listdir(VIDEO_DIR):
        if os.path.splitext(v)[0].lower() == avatar.lower():
            vid = os.path.join(VIDEO_DIR, v)
            break
    if not vid or not audio_path:
        return jsonify(video_url="")
    out = AudioVideoSync().sync_audio_video(vid, audio_path)
    return jsonify(video_url="/" + out)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)