#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import subprocess
import threading
import time
import os
import sys
from datetime import datetime, timedelta

# 配置参数
RTSP_URL = "rtsp://10.68.66.174/audio"
RECORD_DURATION_MINUTES = 60  # 录制时长（分钟）
REST_DURATION_MINUTES = 20    # 休息时长（分钟）
START_HOUR = 8               # 开始时间（小时）
END_HOUR = 20                # 结束时间（小时）
OUTPUT_DIR = "./"            # 输出目录

# 录制状态
recording_state = {
    "is_recording": False,
    "current_file": None,
    "start_time": None,
    "total_recorded": 0
}

def log(msg):
    """带时间戳的日志输出"""
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {msg}")
    sys.stdout.flush()

def check_dependencies():
    """检查依赖是否满足"""
    errors = []
    
    # 检查 Python 版本
    if sys.version_info[0] < 3:
        errors.append("需要 Python 3.x")
    
    # 检查 ffmpeg
    try:
        result = subprocess.run(['ffmpeg', '-version'], capture_output=True, text=True, timeout=5)
        if result.returncode != 0:
            errors.append("ffmpeg 未正确安装")
        else:
            log("ffmpeg 已安装: " + result.stdout.split('\n')[0])
    except FileNotFoundError:
        errors.append("未找到 ffmpeg，请安装: yum install ffmpeg 或 apt install ffmpeg")
    except Exception as e:
        errors.append(f"检查 ffmpeg 时出错: {e}")
    
    return errors

def get_output_path():
    """生成带时间戳的文件名"""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    return os.path.join(OUTPUT_DIR, f"rtsp_{timestamp}.wav")

def is_within_schedule():
    """检查当前是否在有效时间段内"""
    current_hour = datetime.now().hour
    return START_HOUR <= current_hour < END_HOUR

def do_record(output_file):
    """执行 ffmpeg 录制"""
    cmd = [
        'ffmpeg',
        '-rtsp_transport', 'tcp',
        '-i', RTSP_URL,
        '-t', str(RECORD_DURATION_MINUTES * 60),
        '-acodec', 'pcm_s16le',
        '-ar', '16000',
        '-ac', '1',
        '-y',
        output_file
    ]
    
    log(f"执行命令: {' '.join(cmd)}")
    
    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True
    )
    
    # 监控录制进度
    start_time = time.time()
    last_size = 0
    
    while process.poll() is None:
        if os.path.exists(output_file):
            current_size = os.path.getsize(output_file)
            if current_size != last_size:
                size_mb = current_size / (1024 * 1024)
                elapsed = time.time() - start_time
                progress = (elapsed / (RECORD_DURATION_MINUTES * 60)) * 100
                print(f"\r进度: {int(elapsed)}秒 / {RECORD_DURATION_MINUTES * 60}秒 ({progress:.1f}%), 大小: {size_mb:.2f}MB", end='')
                sys.stdout.flush()
                last_size = current_size
        time.sleep(1)
    
    print()  # 换行
    
    # 检查进程返回码
    if process.returncode != 0:
        log(f"ffmpeg 返回错误码: {process.returncode}")
        remaining = process.stdout.read()
        if remaining:
            log(f"ffmpeg 输出: {remaining[:500]}")
    
    return process.returncode, output_file

def record_audio(check_schedule=True):
    """执行音频录制
    check_schedule: 是否检查时间段限制
    """
    if recording_state["is_recording"]:
        log("正在录制中，跳过本次录制")
        return
    
    if check_schedule and not is_within_schedule():
        log(f"不在有效时间段内（{START_HOUR}:00-{END_HOUR}:00），跳过本次录制")
        return
    
    output_file = get_output_path()
    log(f"开始录制: {output_file}")
    
    recording_state["is_recording"] = True
    recording_state["current_file"] = output_file
    recording_state["start_time"] = datetime.now()
    
    try:
        returncode, file_path = do_record(output_file)
        
        # 录制完成
        if os.path.exists(file_path):
            size_mb = os.path.getsize(file_path) / (1024 * 1024)
            recording_state["total_recorded"] += 1
            log(f"录制完成: {file_path} ({size_mb:.2f}MB)")
        else:
            log("录制失败，文件未生成")
            
    except FileNotFoundError as e:
        log(f"命令未找到: {e}")
        log("请确保 ffmpeg 已安装并可在 PATH 中访问")
    except Exception as e:
        log(f"录制错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        recording_state["is_recording"] = False
        recording_state["current_file"] = None
        recording_state["start_time"] = None

def handle_rest():
    """处理休息时间"""
    if is_within_schedule():
        log(f"进入休息模式，等待 {REST_DURATION_MINUTES} 分钟...")
        time.sleep(REST_DURATION_MINUTES * 60)
        log("休息结束")
    else:
        # 计算到下一个有效时间段
        now = datetime.now()
        if now.hour >= END_HOUR:
            next_day = now + timedelta(days=1)
            target = next_day.replace(hour=START_HOUR, minute=0, second=0, microsecond=0)
        else:
            target = now.replace(hour=START_HOUR, minute=0, second=0, microsecond=0)
        
        wait_seconds = (target - now).total_seconds()
        log(f"不在有效时间段，等待 {wait_seconds/3600:.1f} 小时到 {START_HOUR}:00...")
        time.sleep(wait_seconds)

def run_scheduler():
    """运行定时器"""
    log("定时录制服务已启动")
    log(f"  - 录制时长: {RECORD_DURATION_MINUTES} 分钟")
    log(f"  - 休息时长: {REST_DURATION_MINUTES} 分钟")
    log(f"  - 有效时段: {START_HOUR}:00 - {END_HOUR}:00")
    log("-" * 50)
    
    # 启动时立即开始录制（无视时间段）
    log("启动脚本，立即开始录制（本次无视时间段）")
    record_audio(check_schedule=False)
    
    # 录制完成后处理休息
    if recording_state["total_recorded"] > 0:
        handle_rest()
    
    while True:
        now = datetime.now()
        
        # 检查是否在有效时间段内
        if is_within_schedule():
            # 每小时整点开始录制
            if now.minute == 0 and now.second == 0:
                log("整点时间到，开始录制")
                record_audio(check_schedule=True)
        
        time.sleep(1)

def main():
    """主函数"""
    print("=" * 50)
    print("RTSP 音频录制服务")
    print("=" * 50)
    
    # 检查依赖
    log("检查系统依赖...")
    errors = check_dependencies()
    
    if errors:
        log("发现以下问题:")
        for error in errors:
            log(f"  - {error}")
        log("请修复上述问题后重试")
        sys.exit(1)
    
    # 确保输出目录存在
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    log(f"输出目录: {os.path.abspath(OUTPUT_DIR)}")
    
    # 启动定时器
    try:
        run_scheduler()
    except KeyboardInterrupt:
        log("收到停止信号，正在退出...")
        sys.exit(0)

if __name__ == '__main__':
    main()
