main.py

# -*- coding: utf-8 -*-
import atexit
import logging
import os
import time
import traceback
import uuid
from concurrent.futures import ThreadPoolExecutor

from flask import Flask, request, jsonify
from flask_cors import CORS  # Flask的跨域处理组件
from werkzeug.utils import secure_filename

try:
    from asr_funasr import funasr
except Exception as e:
    import sys
    traceback.print_exc(file=sys.stderr)
    raise


# 配置日志
# ===================== 1. 修复日志配置 =====================
# 清除可能存在的日志配置
logging.getLogger().handlers.clear()

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)

# 设置日志格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_handler.setFormatter(formatter)

# 配置根日志器
root_logger = logging.getLogger()
root_logger.setLevel(logging.INFO)
root_logger.addHandler(console_handler)

# 创建模块专用日志器
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


"""
app global settings and configuration
"""

# ===================== 2. 配置管理优化 =====================
class Config:
    # 文件上传配置
    UPLOAD_DIR = os.getenv('UPLOAD_DIRECTORY')
    if not UPLOAD_DIR:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        UPLOAD_DIR = os.path.join(current_dir, "wav_uploads")
    
    # 文件大小限制（默认50MB）
    MAX_FILE_SIZE = int(os.getenv('MAX_FILE_SIZE', 50 * 1024 * 1024))
    
    # 转录超时时间（秒）
    TRANS_TIMEOUT = float(os.getenv('TRANS_TIMEOUT', 120.0))
    
    # 线程池大小
    TRANS_WORKER = int(os.getenv('TRANS_WORKER', 5))
    
    # 最大队列长度
    MAX_QUEUE_SIZE = int(os.getenv('MAX_QUEUE_SIZE', 25))
    
    # 支持的文件格式
    ALLOWED_EXTENSIONS = {'.wav'}
    
    @classmethod
    def validate(cls):
        """验证配置"""
        if cls.TRANS_TIMEOUT <= 0:
            raise ValueError("TRANS_TIMEOUT must be positive")
        if cls.TRANS_WORKER <= 0:
            raise ValueError("TRANS_WORKER must be positive")
        if cls.MAX_FILE_SIZE <= 0:
            raise ValueError("MAX_FILE_SIZE must be positive")

# 初始化配置
config = Config()
config.validate()

# 创建上传目录
os.makedirs(config.UPLOAD_DIR, exist_ok=True)

# Flask应用初始化
app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = config.MAX_FILE_SIZE

CORS(
    app,
    resources={
        r"/transcribe/*": {"origins": "*"},
        r"/system/*": {"origins": "*"}
    },
)


# ===================== 3. 并发控制配置优化 =====================
# 创建固定大小的线程池（全局唯一，避免重复创建）
executor = ThreadPoolExecutor(
    max_workers=config.TRANS_WORKER, 
    thread_name_prefix="transcribe_"
)

# 任务队列监控（线程安全）
import threading
task_stats_lock = threading.Lock()
task_stats = {
    "processing": 0,  # 正在处理的任务数
    "queued": 0,      # 排队中的任务数
    "completed": 0,   # 已完成任务数
    "failed": 0       # 失败任务数
}

# 注册应用关闭时清理资源
@atexit.register
def cleanup_resources():
    """Clean up resources on application shutdown"""
    logger.info("Shutting down thread pool...")
    executor.shutdown(wait=True)
    logger.info("Thread pool shutdown complete")

def update_task_stats(field: str, delta: int = 1):
    """Thread-safe update of task statistics"""
    with task_stats_lock:
        task_stats[field] += delta

def get_task_stats():
    """Thread-safe retrieval of task statistics"""
    with task_stats_lock:
        return task_stats.copy()
# ========================================================


def run_transcription(file_path: str, output_type: str = "txt") -> dict:
    """Execute transcription task synchronously and return results"""
    start_time = time.time()
    try:
        # Check if file exists
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Audio file not found: {file_path}")
        
        # Check file size
        file_size = os.path.getsize(file_path)
        if file_size == 0:
            raise ValueError("Audio file is empty")
        
        logger.info(f"Starting transcription for file: {file_path}, size: {file_size} bytes")
        
        # Call funasr for transcription
        result = funasr.transcribe(file_path, output_type)
        
        processing_time = time.time() - start_time
        logger.info(f"Transcription completed, duration: {processing_time:.2f}s")
        
        update_task_stats("completed")
        
        return {
            "status": "success",
            "transcription": result,
            "processing_time": f"{processing_time:.2f}s",
            "file_size": file_size,
            "output_type": output_type
        }
    except Exception as e:
        processing_time = time.time() - start_time
        error_msg = str(e)
        logger.error(f"Transcription failed: {error_msg}, duration: {processing_time:.2f}s")
        logger.error(traceback.format_exc())
        
        update_task_stats("failed")
        
        return {
            "status": "error",
            "error": error_msg,
            "processing_time": f"{processing_time:.2f}s",
            "error_type": type(e).__name__
        }


def validate_wav_file(file) -> tuple[bool, str]:
    """Validate if file is in WAV format
    
    Returns:
        tuple: (is_valid, error_message)
    """
    try:
        # Check filename
        if not file.filename:
            return False, "Filename is empty"
        
        # Secure filename check
        filename = secure_filename(file.filename)
        if not filename:
            return False, "Invalid filename"
        
        # Check file extension
        file_ext = os.path.splitext(filename)[1].lower()
        if file_ext not in config.ALLOWED_EXTENSIONS:
            return False, f"Unsupported file format: {file_ext}, only supports: {', '.join(config.ALLOWED_EXTENSIONS)}"
        
        # Check file header (WAV files start with "RIFF")
        current_position = file.tell()
        file.seek(0)
        header = file.read(4)
        file.seek(current_position)  # Restore original position
        
        if len(header) < 4:
            return False, "File corrupted or too small"
        
        if header != b'RIFF':
            return False, "Not a valid WAV file format"
        
        return True, ""
        
    except Exception as e:
        logger.error(f"File validation failed: {str(e)}")
        return False, f"Error occurred during file validation: {str(e)}"


@app.route("/transcribe/file", methods=["POST"])
def transcribe_file():
    """Upload WAV file and perform speech recognition"""
    save_path = ""
    file_id = None
    
    try:
        # 1. Validate request
        if 'file' not in request.files:
            return jsonify({
                "status": "error", 
                "detail": "No uploaded file found",
                "error_code": "NO_FILE"
            }), 400

        file = request.files['file']
        
        if file.filename == '':
            return jsonify({
                "status": "error", 
                "detail": "No file selected",
                "error_code": "EMPTY_FILENAME"
            }), 400

        # 2. Validate file format
        is_valid, error_message = validate_wav_file(file)
        if not is_valid:
            return jsonify({
                "status": "error",
                "detail": error_message,
                "error_code": "INVALID_FORMAT"
            }), 400

        # 3. Check current task load
        current_stats = get_task_stats()
        current_load = current_stats["processing"] + current_stats["queued"]
        if current_load >= config.MAX_QUEUE_SIZE:
            return jsonify({
                "status": "error",
                "detail": f"Server busy, please try again later (current queue: {current_load} tasks)",
                "error_code": "SERVER_BUSY",
                "queue_info": current_stats
            }), 503

        # 4. Save file locally
        file_id = str(uuid.uuid4())
        filename = secure_filename(file.filename)
        file_ext = os.path.splitext(filename)[1]
        save_path = os.path.join(config.UPLOAD_DIR, f"{file_id}{file_ext}")
        
        # Ensure directory exists
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        
        file.save(save_path)
        logger.info(f"File uploaded successfully: {filename} -> {save_path}")

        # 5. Get output type parameter
        output_type = request.form.get('output_type', 'txt')
        if output_type not in ['txt', 'srt']:
            output_type = 'txt'

        # 6. Submit transcription task
        update_task_stats("queued")
        logger.info(f"Task queued, current stats: {get_task_stats()}")

        try:
            # Switch task status: from queued to processing
            update_task_stats("queued", -1)
            update_task_stats("processing")

            # Execute task (synchronous mode with timeout)
            future = executor.submit(run_transcription, save_path, output_type)
            transcription_result = future.result(timeout=config.TRANS_TIMEOUT)
        
        except TimeoutError:
            logger.error(f"Task timeout, file: {save_path}")
            update_task_stats("failed")
            return jsonify({
                "status": "error",
                "detail": f"Transcription task timeout ({config.TRANS_TIMEOUT}s)",
                "error_code": "TIMEOUT",
                "file_id": file_id
            }), 504
        
        finally:
            # Update processing task count regardless of success or failure
            update_task_stats("processing", -1)

        # 7. Return results
        response_data = {
            "status": "completed",
            "file_id": file_id,
            "filename": filename,
            "result": transcription_result
        }
        
        if transcription_result["status"] == "error":
            response_data["status"] = "error"
            response_data["detail"] = transcription_result["error"]
            response_data["error_code"] = "TRANSCRIPTION_FAILED"
        
        return jsonify(response_data)
        
    except Exception as e:
        logger.error(f"API processing failed: {str(e)}")
        logger.error(traceback.format_exc())
        
        # Ensure task statistics are correct
        update_task_stats("failed")
        
        return jsonify({
            "status": "error",
            "detail": f"Error occurred during processing: {str(e)}",
            "error_code": "INTERNAL_ERROR",
            "file_id": file_id
        }), 500
    
    finally:
        # Clean up temporary files
        if save_path and os.path.exists(save_path):
            try:
                os.remove(save_path)
                logger.debug(f"Temporary file deleted: {save_path}")
            except Exception as e:
                logger.warning(f"Failed to delete temporary file: {save_path}, error: {str(e)}")


@app.route("/system/status", methods=["GET"])
def get_system_status():
    """Get system status information"""
    current_stats = get_task_stats()
    return jsonify({
        "status": "running",
        "config": {
            "max_workers": config.TRANS_WORKER,
            "max_file_size": config.MAX_FILE_SIZE,
            "trans_timeout": config.TRANS_TIMEOUT,
            "max_queue_size": config.MAX_QUEUE_SIZE,
            "allowed_extensions": list(config.ALLOWED_EXTENSIONS)
        },
        "task_stats": current_stats,
        "load_percentage": round((current_stats["processing"] + current_stats["queued"]) / config.MAX_QUEUE_SIZE * 100, 2),
        "upload_dir": config.UPLOAD_DIR
    })

@app.route("/system/health", methods=["GET"])
def health_check():
    """Health check endpoint"""
    try:
        # Check if model is working properly
        funasr.init_model()
        
        # Check if upload directory is writable
        test_file = os.path.join(config.UPLOAD_DIR, f"test_{uuid.uuid4()}.tmp")
        with open(test_file, 'w') as f:
            f.write("test")
        os.remove(test_file)
        
        return jsonify({
            "status": "healthy",
            "timestamp": time.time(),
            "checks": {
                "model": "ok",
                "upload_dir": "ok",
                "executor": "ok" if not executor._shutdown else "shutdown"
            }
        })
    except Exception as e:
        logger.error(f"Health check failed: {str(e)}")
        return jsonify({
            "status": "unhealthy",
            "timestamp": time.time(),
            "error": str(e)
        }), 503

@app.route("/system/readiness", methods=["GET"])
def readiness_check():
    """Readiness check endpoint - checks if model is loaded and ready"""
    try:
        if funasr.is_ready():
            return jsonify({
                "status": "ready",
                "timestamp": time.time(),
                "model_info": funasr.get_model_info()
            })
        else:
            return jsonify({
                "status": "not_ready",
                "timestamp": time.time(),
                "message": "Model not initialized yet"
            }), 503
    except Exception as e:
        logger.error(f"Readiness check failed: {str(e)}")
        return jsonify({
            "status": "error",
            "timestamp": time.time(),
            "error": str(e)
        }), 503


if __name__ == "__main__":
    # ===================== 4. Application Startup =====================
    import sys
    logger.info("Python version: %s", sys.version)
    
    # Check if lazy loading is enabled
    lazy_load = os.getenv('LAZY_LOAD_MODEL', 'false').lower() == 'true'

    try:
        if not lazy_load:
            # Preload model to avoid timeout on first request
            logger.info("Preloading FunASR model...")
            
            # Check if funasr object exists
            if funasr is None:
                raise RuntimeError("funasr object is None")
                
            funasr.init_model()
            logger.info("FunASR model loaded successfully")
        else:
            logger.info("Lazy loading enabled - model will be loaded on first request")
        
        # Print configuration information
        logger.info("Configuration:")
        logger.info("  - Max worker threads: %d", config.TRANS_WORKER)
        logger.info("  - File size limit: %.1fMB", config.MAX_FILE_SIZE / 1024 / 1024)
        logger.info("  - Transcription timeout: %ds", config.TRANS_TIMEOUT)
        logger.info("  - Max queue length: %d", config.MAX_QUEUE_SIZE)
        logger.info("  - Upload directory: %s", config.UPLOAD_DIR)
        logger.info("  - Lazy load model: %s", lazy_load)
        
        # Start application
        logger.info("Starting Flask application...")
        
        # Force flush output
        sys.stdout.flush()
        sys.stderr.flush()
        
        app.run(
            host="0.0.0.0", 
            port=int(os.getenv('PORT', 8000)),
            debug=False,  # Don't enable debug in production
            threaded=True
        )
        
    except KeyboardInterrupt:
        logger.info("Received interrupt signal, shutting down application...")
    except Exception as e:
        logger.error("Application startup failed: %s", str(e))
        logger.error(traceback.format_exc())
        # Print more detailed error information
        print(f"Application error details: {traceback.format_exc()}", file=sys.stderr)
    finally:
        logger.info("Application shutdown complete")