# -*- coding: utf-8 -*-
"""
CCA 唱词助手 — Flask 路由
POST /api/cca/upload     上传 A稿+音频，启动流水线
GET  /api/cca/status/<id> 轮询任务状态
GET  /api/cca/review/<id> 获取审稿数据
POST /api/cca/save/<id>   保存编导修改
POST /api/cca/generate/<id> 生成最终 SRT
GET  /api/cca/download/<id> 下载 SRT zip
"""

import json
import os
import sys
import uuid
import threading
import traceback
import zipfile
from io import BytesIO
from pathlib import Path
from datetime import datetime

from flask import Blueprint, request, jsonify, send_file

bp = Blueprint('cca', __name__, url_prefix='/api/cca')

# 运行时数据目录
CCA_DATA_DIR = Path('/workspace/military_tech_voice/backend/cca_data')
CCA_DATA_DIR.mkdir(parents=True, exist_ok=True)

# CCA 源码目录
CCA_SRC_DIR = Path('/workspace/military_tech_voice/backend/cca_src')
if str(CCA_SRC_DIR) not in sys.path:
    sys.path.insert(0, str(CCA_SRC_DIR))

# 任务状态存储（内存，重启丢失无所谓——编导重新上传即可）
tasks = {}


def _get_task(task_id):
    t = tasks.get(task_id)
    if not t:
        return None
    return t


@bp.route('/upload', methods=['POST'])
def upload():
    """接收 A稿 + 音频，创建任务"""
    if 'audio' not in request.files:
        return jsonify({'error': '请上传音频文件'}), 400

    audio_file = request.files['audio']
    script_file = request.files.get('script')

    if not audio_file.filename:
        return jsonify({'error': '音频文件为空'}), 400

    task_id = datetime.now().strftime('%Y%m%d_%H%M%S') + '_' + uuid.uuid4().hex[:6]
    task_dir = CCA_DATA_DIR / task_id
    task_dir.mkdir(parents=True, exist_ok=True)

    # 保存音频
    audio_ext = os.path.splitext(audio_file.filename)[1] or '.mp3'
    audio_path = task_dir / f'audio{audio_ext}'
    audio_file.save(str(audio_path))

    # 保存 A稿
    script_path = None
    if script_file and script_file.filename:
        script_ext = os.path.splitext(script_file.filename)[1] or '.docx'
        script_path = task_dir / f'script{script_ext}'
        script_file.save(str(script_path))

    tasks[task_id] = {
        'id': task_id,
        'status': 'uploaded',
        'progress': '文件已上传，准备处理...',
        'audio_path': str(audio_path),
        'script_path': str(script_path) if script_path else None,
        'created_at': datetime.now().isoformat(),
        'error': None,
        'asr_sentences': None,
        'proofread_sentences': None,
        'review_data': None,
        'final_srt_dir': None,
    }

    # 启动后台处理
    thread = threading.Thread(target=_run_pipeline, args=(task_id,), daemon=True)
    thread.start()

    return jsonify({'task_id': task_id, 'status': 'processing'})


def _run_pipeline(task_id):
    """后台运行 CCA 流水线"""
    task = tasks[task_id]
    try:
        task['status'] = 'processing'
        task['progress'] = '正在提取热词...'

        audio_path = task['audio_path']
        script_path = task['script_path']

        # WAV/大文件 → MP3 压缩（讯飞上传大文件容易超时）
        if audio_path.lower().endswith('.wav'):
            import subprocess
            mp3_path = audio_path.rsplit('.', 1)[0] + '.mp3'
            task['progress'] = '正在压缩音频（WAV→MP3）...'
            subprocess.run(
                ['ffmpeg', '-i', audio_path, '-b:a', '128k', '-y', mp3_path],
                capture_output=True, timeout=300,
            )
            if os.path.exists(mp3_path) and os.path.getsize(mp3_path) > 0:
                audio_path = mp3_path
                task['audio_path'] = mp3_path

        from hotword_extractor import extract_hotwords, read_docx_text, read_text_file
        from asr_client import transcribe, parse_result
        from term_normalizer import normalize_terms
        from ai_proofreader import proofread_batch
        from segment_splitter import split_into_segments

        # Step 0: 热词提取
        hot_words = []
        script_text = ""
        if script_path:
            hot_words = extract_hotwords(script_path, use_ai=True)
            ext = os.path.splitext(script_path)[1].lower()
            if ext == '.docx':
                script_text = read_docx_text(script_path)
            else:
                script_text = read_text_file(script_path)

        task['progress'] = f'热词提取完成({len(hot_words)}个)，正在 ASR 转写...'

        # Step 1: ASR
        sentences, raw_json = transcribe(audio_path, hot_words=hot_words if hot_words else None)

        # 缓存 ASR 原始结果
        task_dir = Path(audio_path).parent
        cache_path = task_dir / 'asr_raw.json'
        with open(cache_path, 'w', encoding='utf-8') as f:
            f.write(raw_json)

        task['progress'] = f'ASR 完成({len(sentences)}句)，正在术语格式化...'

        # 保存 ASR 原始句子（校对前，供 diff 对比）
        asr_original = [(bg, ed, text, spk) for bg, ed, text, spk in sentences]

        # Step 1.5: 术语格式化
        if script_text:
            sentences = normalize_terms(sentences, script_text)

        task['progress'] = '术语格式化完成，正在 AI 校对...'

        # Step 2: AI 校对
        if script_text:
            sentences = proofread_batch(sentences, script_text)

        # Step 2.5: 校对后二次正则
        if script_text:
            sentences = normalize_terms(sentences, script_text)

        task['progress'] = 'AI 校对完成，正在准备审稿数据...'

        # 保存校对后的句子
        task['asr_sentences'] = asr_original
        task['proofread_sentences'] = sentences

        # 构建审稿数据：逐句对比
        review_items = []
        for i, ((bg, ed, orig_text, spk), (_, _, proof_text, _)) in enumerate(
            zip(asr_original, sentences)
        ):
            has_change = orig_text != proof_text
            review_items.append({
                'index': i,
                'start_ms': bg,
                'end_ms': ed,
                'speaker_id': spk,
                'original': orig_text,
                'corrected': proof_text,
                'edited': proof_text,
                'has_change': has_change,
                'confirmed': not has_change,
            })

        task['review_data'] = review_items
        task['status'] = 'review'
        task['progress'] = '审稿数据就绪，请编导审阅确认'

        # 同时保存到磁盘（防丢）
        review_path = task_dir / 'review_data.json'
        with open(review_path, 'w', encoding='utf-8') as f:
            json.dump(review_items, f, ensure_ascii=False, indent=2)

    except Exception as e:
        task['status'] = 'error'
        err_msg = str(e)
        if '余额不足' in err_msg or 'insufficient' in err_msg.lower() or '10317' in err_msg:
            task['error'] = '讯飞录音文件转写余额不足，请联系管理员充值'
        else:
            task['error'] = f'处理出错: {err_msg}'
        task['progress'] = task['error']
        traceback.print_exc()


@bp.route('/status/<task_id>', methods=['GET'])
def status(task_id):
    task = _get_task(task_id)
    if not task:
        return jsonify({'error': '任务不存在'}), 404
    return jsonify({
        'task_id': task_id,
        'status': task['status'],
        'progress': task['progress'],
        'error': task['error'],
    })


@bp.route('/review/<task_id>', methods=['GET'])
def review(task_id):
    task = _get_task(task_id)
    if not task:
        return jsonify({'error': '任务不存在'}), 404
    if task['status'] not in ('review', 'completed'):
        return jsonify({'error': '任务尚未就绪', 'status': task['status']}), 400
    return jsonify({
        'task_id': task_id,
        'items': task['review_data'],
    })


@bp.route('/save/<task_id>', methods=['POST'])
def save(task_id):
    """保存编导的修改（自动保存用）"""
    task = _get_task(task_id)
    if not task:
        return jsonify({'error': '任务不存在'}), 404

    data = request.get_json()
    edits = data.get('items', [])

    for edit in edits:
        idx = edit.get('index')
        if idx is not None and 0 <= idx < len(task['review_data']):
            task['review_data'][idx]['edited'] = edit.get('edited', task['review_data'][idx]['edited'])
            task['review_data'][idx]['confirmed'] = edit.get('confirmed', True)

    return jsonify({'ok': True})


@bp.route('/generate/<task_id>', methods=['POST'])
def generate(task_id):
    """用编导确认后的文本生成最终 SRT"""
    task = _get_task(task_id)
    if not task:
        return jsonify({'error': '任务不存在'}), 404
    if task['status'] not in ('review', 'completed'):
        return jsonify({'error': '任务状态不对'}), 400

    try:
        from ai_line_breaker import process_sentences_with_ai
        from srt_writer import write_srt, ms_to_srt_time
        from segment_splitter import split_into_segments

        # 用编导确认后的文本重建句子列表
        confirmed_sentences = []
        for item in task['review_data']:
            text = item['edited']
            confirmed_sentences.append((
                item['start_ms'], item['end_ms'], text, item['speaker_id']
            ))

        # 切分节目结构
        segments = split_into_segments(confirmed_sentences)

        # 折行 + 生成 SRT
        task_dir = Path(task['audio_path']).parent
        srt_dir = task_dir / 'srt_output'
        srt_dir.mkdir(exist_ok=True)

        srt_files = []
        for seg_name, seg_sentences in segments:
            if not seg_sentences:
                continue
            subtitle_lines = process_sentences_with_ai(seg_sentences)
            srt_path = srt_dir / f'{seg_name}.srt'
            write_srt(subtitle_lines, str(srt_path))
            srt_files.append(str(srt_path))

        task['final_srt_dir'] = str(srt_dir)
        task['status'] = 'completed'
        task['progress'] = f'生成完成，共 {len(srt_files)} 个 SRT 文件'

        return jsonify({
            'ok': True,
            'srt_count': len(srt_files),
            'message': task['progress'],
        })

    except Exception as e:
        return jsonify({'error': f'生成 SRT 出错: {str(e)}'}), 500


@bp.route('/download/<task_id>', methods=['GET'])
def download(task_id):
    """下载 SRT zip 包"""
    task = _get_task(task_id)
    if not task:
        return jsonify({'error': '任务不存在'}), 404
    if not task.get('final_srt_dir'):
        return jsonify({'error': 'SRT 尚未生成'}), 400

    srt_dir = Path(task['final_srt_dir'])
    srt_files = sorted(srt_dir.glob('*.srt'))
    if not srt_files:
        return jsonify({'error': '无 SRT 文件'}), 404

    # 打包 zip
    buf = BytesIO()
    with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
        for srt_file in srt_files:
            zf.write(srt_file, srt_file.name)
    buf.seek(0)

    filename = f'唱词字幕_{task_id}.zip'
    return send_file(buf, mimetype='application/zip', as_attachment=True, download_name=filename)