import os
from flask import Flask, render_template, request, jsonify, send_file
from transformers import AutoModelForCausalLM, AutoTokenizer
from deep_translator import GoogleTranslator
from gtts import gTTS
import uuid

# Initialize Flask app
app = Flask(__name__)

# Load the DeepSeek-V3 model and tokenizer
model_name = "deepseek-ai/DeepSeek-V3"
revision = "main"  # Pin to a specific revision if needed

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, revision=revision)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        revision=revision,
        quantization_config=None  # Disable quantization
    )
except Exception as e:
    print(f"Error loading model: {e}")
    raise

# Supported languages for translation
SUPPORTED_LANGUAGES = {
    "english": "en", "hindi": "hi", "bengali": "bn", "telugu": "te", "marathi": "mr",
    "tamil": "ta", "gujarati": "gu", "kannada": "kn", "malayalam": "ml", "punjabi": "pa",
    "odia": "or", "urdu": "ur", "assamese": "as", "sanskrit": "sa", "arabic": "ar",
    "chinese": "zh-CN", "dutch": "nl", "french": "fr", "filipino": "tl", "greek": "el",
    "indonesian": "id", "italian": "it", "japanese": "ja", "korean": "ko", "latin": "la",
    "nepali": "ne", "portuguese": "pt", "romanian": "ro", "russian": "ru", "spanish": "es",
    "swedish": "sv", "thai": "th", "ukrainian": "uk", "turkish": "tr"
}

# Temporary directory for audio files
AUDIO_DIR = "audio"
os.makedirs(AUDIO_DIR, exist_ok=True)

# Handle user queries using DeepSeek-V3
def handle_query(query):
    inputs = tokenizer(query, return_tensors="pt")
    outputs = model.generate(
        inputs.input_ids,
        max_length=50,
        num_return_sequences=1,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.strip()

# Generate response and translate if needed
def generate_response(query, language):
    try:
        bot_response = handle_query(query)
        target_lang = SUPPORTED_LANGUAGES.get(language.lower(), "en")
        if target_lang != "en":
            bot_response = GoogleTranslator(source='en', target=target_lang).translate(bot_response)
        return bot_response
    except Exception as e:
        return f"Error fetching the response: {str(e)}"

# Convert text to speech using gTTS
def text_to_speech(text, lang="en"):
    try:
        audio_filename = f"{uuid.uuid4()}.mp3"
        audio_path = os.path.join(AUDIO_DIR, audio_filename)
        tts = gTTS(text=text, lang=lang)
        tts.save(audio_path)
        return audio_path
    except Exception as e:
        print(f"Error generating speech: {e}")
        return None

# Flask routes
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/chat', methods=['POST'])
def chat():
    try:
        data = request.json
        user_message = data.get("message")
        language = data.get("language", "english").lower()

        if not user_message:
            return jsonify({"response": "Please say something!", "audio_url": None})

        bot_response = generate_response(user_message, language)
        target_lang = SUPPORTED_LANGUAGES.get(language, "en")
        audio_path = text_to_speech(bot_response, lang=target_lang)

        if audio_path:
            audio_url = f"/audio/{os.path.basename(audio_path)}"
        else:
            audio_url = None

        return jsonify({"response": bot_response, "audio_url": audio_url})
    except Exception as e:
        return jsonify({"response": f"An error occurred: {str(e)}", "audio_url": None})

@app.route('/audio/<filename>')
def serve_audio(filename):
    return send_file(os.path.join(AUDIO_DIR, filename), mimetype="audio/mp3")

if __name__ == '__main__':
    app.run(debug=True)