import os
from flask import Flask, request, jsonify
from flask_swagger_ui import get_swaggerui_blueprint
from transformers import pipeline
import torch
from pydub import AudioSegment
from moviepy.editor import VideoFileClip


class Config:
    MODEL_NAME = "openai/whisper-large-v3-turbo"
    BATCH_SIZE = 8
    DEVICE = 0 if torch.cuda.is_available() else "cpu"
    TEMP_DIR = './temp_files'

    # Swagger configuration
    SWAGGER_URL = '/swagger'
    API_URL = '/static/swagger.json'


class TranscriptionService:
    def __init__(self, model_name=Config.MODEL_NAME):
        self.pipeline = pipeline(
            task="automatic-speech-recognition",
            model=model_name,
            chunk_length_s=30,
            device=Config.DEVICE,
        )
        os.makedirs(Config.TEMP_DIR, exist_ok=True)

    def transcribe_audio(self, audio_path, task="transcribe"):
        try:
            result = self.pipeline(audio_path, batch_size=Config.BATCH_SIZE, generate_kwargs={"task": task},
                                   return_timestamps=True)
            return result['text']
        except Exception as exc:
            raise RuntimeError(f"Error during transcription: {exc}")

    @staticmethod
    def save_file(file):
        temp_file_path = os.path.join(Config.TEMP_DIR, file.filename)
        file.save(temp_file_path)
        # Проверяем расширение файла и конвертируем, если необходимо
        file_ext = os.path.splitext(temp_file_path)[1].lower()
        if file_ext == '.m4a':
            wav_file_path = os.path.splitext(temp_file_path)[0] + '.wav'
            audio = AudioSegment.from_file(temp_file_path, format='m4a')
            audio.export(wav_file_path, format='wav')
            os.remove(temp_file_path)  # Удаляем исходный .m4a файл
            return wav_file_path
        return temp_file_path

    @staticmethod
    def extract_audio_from_video(video_path):
        try:
            video = VideoFileClip(video_path)
            audio_path = os.path.splitext(video_path)[0] + '.wav'
            video.audio.write_audiofile(audio_path, codec='pcm_s16le')
            return audio_path
        except Exception as exc:
            raise RuntimeError(f"Error extracting audio from video: {exc}")

    @staticmethod
    def delete_file(file_path):
        if os.path.exists(file_path):
            os.remove(file_path)


app = Flask(__name__)
transcription_service = TranscriptionService()


@app.route('/transcribe', methods=['POST'])
def transcribe_audio_route():
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400

    file = request.files['file']
    file_path = transcription_service.save_file(file)

    try:
        transcription = transcription_service.transcribe_audio(file_path)
        return jsonify({'transcription': transcription})
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        transcription_service.delete_file(file_path)


@app.route('/transcribe_video', methods=['POST'])
def transcribe_video_route():
    if 'file' not in request.files:
        return jsonify({'error': 'No video file uploaded'}), 400

    file = request.files['file']
    video_file_path = os.path.join(Config.TEMP_DIR, file.filename)
    file.save(video_file_path)

    # Проверяем расширение файла
    file_ext = os.path.splitext(video_file_path)[1].lower()
    supported_video_formats = ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv']

    if file_ext not in supported_video_formats:
        transcription_service.delete_file(video_file_path)
        return jsonify({'error': f'Unsupported video format: {file_ext}'}), 400

    try:
        # Извлекаем аудио из видео
        audio_path = transcription_service.extract_audio_from_video(video_file_path)

        # Транскрибируем аудио
        transcription = transcription_service.transcribe_audio(audio_path)

        return jsonify({'transcription': transcription})
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        # Удаляем как видеофайл, так и извлеченное аудио
        transcription_service.delete_file(video_file_path)
        if 'audio_path' in locals():
            transcription_service.delete_file(audio_path)


# Настройка Swagger
swaggerui_blueprint = get_swaggerui_blueprint(
    Config.SWAGGER_URL,
    Config.API_URL,
    config={
        'app_name': "Audio and Video Transcription API"
    }
)

app.register_blueprint(swaggerui_blueprint, url_prefix=Config.SWAGGER_URL)

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

Почему у меня транскрибируется только концовка длинного видео?