import gradio as gr
from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
import os
import uuid
from werkzeug.utils import secure_filename
import json
from datetime import datetime

from services.translator import TranslationService
from services.document_processor import DocumentProcessor
from services.audio_generator import AudioGenerator

app = Flask(__name__, static_folder='static')
CORS(app)

app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['OUTPUT_FOLDER'] = 'outputs'
app.config['ALLOWED_EXTENSIONS'] = {'txt', 'pdf', 'docx', 'md'}

os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['OUTPUT_FOLDER'], exist_ok=True)

translator = TranslationService()
doc_processor = DocumentProcessor()
audio_gen = AudioGenerator()

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

@app.route('/')
def index():
    return app.send_static_file('index.html')

@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({'status': 'healthy', 'timestamp': datetime.now().isoformat()})

@app.route('/languages', methods=['GET'])
def get_supported_languages():
    languages = {
        'en-hi': 'English to Hindi',
        'hi-en': 'Hindi to English', 
        'en-fr': 'English to French',
        'en-es': 'English to Spanish',
        'en-de': 'English to German',
        'en-zh': 'English to Chinese',
        'en-ja': 'English to Japanese',
        'en-ar': 'English to Arabic'
    }
    return jsonify({'languages': languages})

@app.route('/upload', methods=['POST'])
def upload_document():
    if 'file' not in request.files:
        return jsonify({'error': 'No file provided'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400
    
    if file and allowed_file(file.filename):
        task_id = str(uuid.uuid4())
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{task_id}_{filename}")
        file.save(filepath)
        
        file_extension = filename.rsplit('.', 1)[1].lower()
        
        try:
            extracted_text = doc_processor.extract_text(filepath, file_extension)
            
            return jsonify({
                'task_id': task_id,
                'filename': filename,
                'text_preview': extracted_text[:500] + '...' if len(extracted_text) > 500 else extracted_text,
                'total_characters': len(extracted_text),
                'status': 'uploaded'
            })
        except Exception as e:
            return jsonify({'error': f'Failed to process document: {str(e)}'}), 500
    
    return jsonify({'error': 'Invalid file type'}), 400

@app.route('/translate/<task_id>', methods=['POST'])
def translate_document(task_id):
    data = request.get_json()
    target_language = data.get('target_language', 'en-hi')
    
    try:
        upload_files = [f for f in os.listdir(app.config['UPLOAD_FOLDER']) if f.startswith(task_id)]
        if not upload_files:
            return jsonify({'error': 'Task not found'}), 404
        
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], upload_files[0])
        file_extension = upload_files[0].rsplit('.', 1)[1].lower()
        
        text = doc_processor.extract_text(filepath, file_extension)
        
        chunks = doc_processor.split_into_chunks(text)
        
        translated_chunks = []
        for i, chunk in enumerate(chunks):
            translated = translator.translate(chunk, target_language)
            translated_chunks.append(translated)
        
        translated_text = ' '.join(translated_chunks)
        
        output_path = os.path.join(app.config['OUTPUT_FOLDER'], f"{task_id}_translated.txt")
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(translated_text)
        
        return jsonify({
            'task_id': task_id,
            'status': 'completed',
            'translated_text': translated_text,
            'chunks_processed': len(chunks)
        })
        
    except Exception as e:
        return jsonify({'error': f'Translation failed: {str(e)}'}), 500

@app.route('/generate_audio/<task_id>', methods=['POST'])
def generate_audio(task_id):
    data = request.get_json()
    language_code = data.get('language_code', 'hi')
    
    try:
        translation_file = os.path.join(app.config['OUTPUT_FOLDER'], f"{task_id}_translated.txt")
        if not os.path.exists(translation_file):
            return jsonify({'error': 'Translation not found. Please translate the document first.'}), 404
        
        with open(translation_file, 'r', encoding='utf-8') as f:
            translated_text = f.read()
        
        audio_path = audio_gen.generate_audio(translated_text, task_id, language_code)
        
        return jsonify({
            'task_id': task_id,
            'status': 'completed',
            'audio_path': audio_path,
            'message': 'Audio generated successfully'
        })
        
    except Exception as e:
        return jsonify({'error': f'Audio generation failed: {str(e)}'}), 500

@app.route('/download/<task_id>/<file_type>', methods=['GET'])
def download_file(task_id, file_type):
    try:
        if file_type == 'translation':
            filepath = os.path.join(app.config['OUTPUT_FOLDER'], f"{task_id}_translated.txt")
            if os.path.exists(filepath):
                return send_file(filepath, as_attachment=True, download_name='translation.txt')
        elif file_type == 'audio':
            filepath = os.path.join(app.config['OUTPUT_FOLDER'], f"{task_id}_audio.mp3")
            if os.path.exists(filepath):
                return send_file(filepath, as_attachment=True, download_name='translation_audio.mp3')
        
        return jsonify({'error': 'File not found'}), 404
        
    except Exception as e:
        return jsonify({'error': f'Download failed: {str(e)}'}), 500

# Gradio interface for Hugging Face Spaces
def translate_and_generate_audio(file, target_language):
    if file is None:
        return "Please upload a file", None
    
    # Save uploaded file
    task_id = str(uuid.uuid4())
    filename = os.path.basename(file.name)
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], f"{task_id}_{filename}")
    
    with open(file.name, 'rb') as f:
        with open(filepath, 'wb') as out:
            out.write(f.read())
    
    file_extension = filename.rsplit('.', 1)[1].lower()
    
    try:
        # Extract text
        text = doc_processor.extract_text(filepath, file_extension)
        
        # Translate
        chunks = doc_processor.split_into_chunks(text)
        translated_chunks = []
        for chunk in chunks:
            translated = translator.translate(chunk, target_language)
            translated_chunks.append(translated)
        
        translated_text = ' '.join(translated_chunks)
        
        # Generate audio
        language_code = target_language.split('-')[1]
        audio_path = audio_gen.generate_audio(translated_text, task_id, language_code)
        
        return translated_text, audio_path
        
    except Exception as e:
        return f"Error: {str(e)}", None

# Create Gradio interface
iface = gr.Interface(
    fn=translate_and_generate_audio,
    inputs=[
        gr.File(label="Upload Document", file_types=[".txt", ".pdf", ".docx", ".md"]),
        gr.Dropdown(
            choices=[
                ("English to Hindi", "en-hi"),
                ("Hindi to English", "hi-en"),
                ("English to French", "en-fr"),
                ("English to Spanish", "en-es"),
                ("English to German", "en-de"),
                ("English to Chinese", "en-zh"),
                ("English to Japanese", "en-ja"),
                ("English to Arabic", "en-ar")
            ],
            label="Target Language",
            value="en-hi"
        )
    ],
    outputs=[
        gr.Textbox(label="Translated Text", lines=10),
        gr.Audio(label="Audio Output", type="filepath")
    ],
    title="Document Translator with Audio Generation",
    description="Upload a document to translate it and generate audio output"
)

if __name__ == '__main__':
    # For Hugging Face Spaces
    if os.environ.get('SPACE_ID'):
        iface.launch()
    else:
        # For local development
        app.run(debug=True, port=5001)