# -*- coding: utf-8 -*-
"""
ai-translate AI 翻译蓝图
支持 NLLB-200 本地翻译引擎
"""
import os
import re
import json
import uuid
import threading
import tempfile
from flask import Blueprint, render_template, request, jsonify, send_file

bp = Blueprint('ai_translate', __name__, url_prefix='/ai-translate')

try:
    from config import BASE_DIR
except ImportError:
    BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

CONFIG_PATH = os.path.join(BASE_DIR, 'config', 'ai_translate_config.json')

# 引擎状态
_model = None
_model_lock = threading.Lock()
_engine_status = {'loaded': False, 'loading': False, 'engine_type': None, 'model_name': None, 'error': None}

_tasks = {}

SUBTITLE_EXTS = ('.srt', '.ass', '.ssa', '.vtt')
LANGUAGES = {
    'zh': '中文', 'en': '英文', 'ja': '日文', 'ko': '韩文',
    'fr': '法文', 'de': '德文', 'es': '西班牙文', 'ru': '俄文',
 'th': '泰文', 'vi': '越南文', 'it': '意大利文',
    'pt': '葡萄牙文', 'auto': '自动检测',
}

# NLLB-200 语言代码映射
_NLLB_LANG_MAP = {
    'zh': 'zho_Hans', 'en': 'eng_Latn', 'ja': 'jpn_Jpan', 'ko': 'kor_Hang',
    'fr': 'fra_Latn', 'de': 'deu_Latn', 'es': 'spa_Latn', 'ru': 'rus_Cyrl',
    'th': 'tha_Thai', 'vi': 'vie_Latn', 'it': 'ita_Latn',
    'pt': 'por_Latn',
}

# langdetect 返回值 → 我们的语言代码
_DETECT_LANG_MAP = {
    'zh-cn': 'zh', 'zh-tw': 'zh', 'zh': 'zh',
    'en': 'en', 'ja': 'ja', 'ko': 'ko',
    'fr': 'fr', 'de': 'de', 'es': 'es', 'ru': 'ru',
    'th': 'th', 'vi': 'vi', 'it': 'it', 'pt': 'pt',
}

ENGINE_TYPES = {
    'nllb': {'name': 'NLLB-200', 'desc': 'HF 目录或模型 ID'},
}


def _default_config():
    return {
        'engine_type': 'nllb',
        'nllb_dir': '',
        'max_new_tokens': 512,
        'batch_size': 8,
        'last_nllb_model': '',
        'last_src_lang': 'en',
        'last_tgt_lang': 'zh',
    }


def _load_config():
    cfg = _default_config()
    if os.path.exists(CONFIG_PATH):
        try:
            with open(CONFIG_PATH, 'r', encoding='utf-8') as f:
                cfg.update(json.load(f))
        except Exception:
            pass
    return cfg


def _save_config(cfg):
    os.makedirs(os.path.dirname(CONFIG_PATH), exist_ok=True)
    with open(CONFIG_PATH, 'w', encoding='utf-8') as f:
        json.dump(cfg, f, ensure_ascii=False, indent=2)



# ==================== 字幕解析 ====================

def _parse_srt(text):
    blocks = re.split(r'\n\s*\n', text.strip())
    entries = []
    for block in blocks:
        lines = block.strip().split('\n')
        if len(lines) < 3:
            continue
        try:
            idx = int(lines[0].strip())
        except ValueError:
            continue
        timecode = lines[1].strip()
        content = '\n'.join(lines[2:]).strip()
        if content:
            entries.append((idx, timecode, content))
    return entries


def _build_srt(entries):
    parts = []
    for idx, tc, content in entries:
        parts.append(f'{idx}\n{tc}\n{content}')
    return '\n\n'.join(parts) + '\n'


def _parse_vtt(text):
    text = re.sub(r'^WEBVTT\s*\n', '', text.strip(), count=1)
    blocks = re.split(r'\n\s*\n', text.strip())
    entries = []
    for i, block in enumerate(blocks):
        lines = block.strip().split('\n')
        if len(lines) < 2:
            continue
        timecode = None
        content_lines = []
        for line in lines:
            if '-->' in line:
                timecode = line.strip()
            elif timecode is not None:
                content_lines.append(line)
            else:
                if '-->' not in line:
                    continue
        if timecode and content_lines:
            entries.append((i + 1, timecode, '\n'.join(content_lines).strip()))
    return entries


def _build_vtt(entries):
    parts = ['WEBVTT\n']
    for idx, tc, content in entries:
        parts.append(f'{idx}\n{tc}\n{content}')
    return '\n\n'.join(parts) + '\n'


def _parse_ass(text):
    entries = []
    for line in text.split('\n'):
        line = line.strip()
        if line.lower().startswith('dialogue:'):
            entries.append(line)
    return entries


def _ass_extract_text(dialogue_line):
    parts = dialogue_line.split(',', 9)
    if len(parts) < 10:
        return dialogue_line, ''
    text = parts[9]
    text = re.sub(r'\{[^}]*\}', '', text)
    text = text.replace('\\N', '\n').replace('\\n', '\n')
    return dialogue_line, text.strip()


def _ass_replace_text(dialogue_line, new_text):
    parts = dialogue_line.split(',', 9)
    if len(parts) < 10:
        return dialogue_line
    translated = new_text.replace('\n', '\\N')
    parts[9] = translated
    return ','.join(parts)


def _build_ass(original_lines, translated_texts):
    result = []
    trans_idx = 0
    for line in original_lines:
        if line.lower().startswith('dialogue:'):
            if trans_idx < len(translated_texts):
                result.append(_ass_replace_text(line, translated_texts[trans_idx]))
                trans_idx += 1
            else:
                result.append(line)
        else:
            result.append(line)
    return '\n'.join(result)


def _detect_subtitle_format(filename, content):
    ext = os.path.splitext(filename)[1].lower()
    if ext == '.srt':
        return 'srt'
    if ext == '.vtt':
        return 'vtt'
    if ext in ('.ass', '.ssa'):
        return 'ass'
    if 'WEBVTT' in content[:20]:
        return 'vtt'
    if 'Dialogue:' in content:
        return 'ass'
    return 'srt'


# ==================== 模型加载与翻译 ====================

def _load_model(model_path, model_name):
    global _model
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
    import torch

    _engine_status['loading'] = True
    _engine_status['error'] = None
    try:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        tokenizer = AutoTokenizer.from_pretrained(model_path, clean_up_tokenization_spaces=True)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        model = model.to(device)
        model.eval()
        _model = {'model': model, 'tokenizer': tokenizer, 'device': device}

        _engine_status['loaded'] = True
        _engine_status['model_name'] = model_name
        _engine_status['engine_type'] = 'nllb'
        _engine_status['error'] = None
        return True, None

    except Exception as e:
        _engine_status['error'] = str(e)
        return False, str(e)
    finally:
        _engine_status['loading'] = False


def _unload_model():
    global _model
    _model = None
    _engine_status['loaded'] = False
    _engine_status['model_name'] = None
    _engine_status['engine_type'] = None
    import gc
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _detect_lang(text):
    """用 langdetect 检测语言，返回我们的语言代码（如 'zh'、'en'）"""
    try:
        from langdetect import detect
        code = detect(text)
        return _DETECT_LANG_MAP.get(code, _DETECT_LANG_MAP.get(code.split('-')[0], 'en'))
    except Exception:
        return 'en'


def _do_translate(texts, src_lang, tgt_lang, cfg):
    if not _model:
        return None, '模型未加载'
    try:
        import torch
        model = _model['model']
        tokenizer = _model['tokenizer']
        device = _model['device']

        nllb_tgt = _NLLB_LANG_MAP.get(tgt_lang, _NLLB_LANG_MAP.get('zh', 'zho_Hans'))
        tgt_token_id = tokenizer.convert_tokens_to_ids(nllb_tgt)

        # 自动检测时，用第一条文本检测语言
        if src_lang == 'auto':
            sample = ' '.join(texts[:3])
            detected = _detect_lang(sample)
            nllb_src = _NLLB_LANG_MAP.get(detected, 'eng_Latn')
        else:
            nllb_src = _NLLB_LANG_MAP.get(src_lang, 'eng_Latn')

        results = []
        for text in texts:
            tokenizer.src_lang = nllb_src
            inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)
            with torch.no_grad():
                translated = model.generate(
                    **inputs,
                    forced_bos_token_id=tgt_token_id,
                    max_new_tokens=cfg.get('max_new_tokens', 512)
                )
            result = tokenizer.decode(translated[0], skip_special_tokens=True)
            results.append(result)
        return results, None
    except Exception as e:
        return None, str(e)


def _ensure_model(src_lang=None, tgt_lang=None):
    if _engine_status['loaded'] and _model:
        return True, None
    cfg = _load_config()
    models_dir = cfg.get('nllb_dir', '')
    if not models_dir or not os.path.isdir(models_dir):
        return False, '请在设置中配置 NLLB 模型目录'

    # 用 last 模型或第一个可用
    model_name = cfg.get('last_nllb_model', '')
    if model_name:
        full = os.path.join(models_dir, model_name)
        if not (os.path.isdir(full) and os.path.exists(os.path.join(full, 'config.json'))):
            model_name = ''
    if not model_name:
        for name in sorted(os.listdir(models_dir)):
            full = os.path.join(models_dir, name)
            if os.path.isdir(full) and os.path.exists(os.path.join(full, 'config.json')):
                model_name = name
                break
    if not model_name:
        return False, '未找到 NLLB 模型，请在设置中配置目录'

    model_path = os.path.join(models_dir, model_name)
    with _model_lock:
        if _engine_status['loaded'] and _model:
            return True, None
        ok, load_err = _load_model(model_path, model_name)
        if ok:
            cfg['last_nllb_model'] = model_name
            _save_config(cfg)
        return ok, load_err


# ==================== 启动时自动加载模型 ====================

_auto_load_done = False


def _trigger_auto_load():
    global _auto_load_done
    if _auto_load_done:
        return
    _auto_load_done = True

    threading.Thread(target=_ensure_model, daemon=True).start()


# ==================== 路由 ====================

@bp.route('/')
def page():
    _trigger_auto_load()
    cfg = _load_config()
    return render_template('ai_translate.html', config=cfg, languages=LANGUAGES, engine_types=ENGINE_TYPES)


@bp.route('/model-status')
def model_status():
    return jsonify({
        'loaded': _engine_status['loaded'],
        'loading': _engine_status['loading'],
        'model_name': _engine_status.get('model_name'),
    })


@bp.route('/config', methods=['GET'])
def get_config():
    return jsonify(_load_config())


@bp.route('/config', methods=['POST'])
def save_config():
    data = request.get_json()
    cfg = _load_config()
    for key in ('nllb_dir', 'last_src_lang', 'last_tgt_lang', 'max_new_tokens', 'batch_size'):
        if key in data:
            cfg[key] = data[key]
    _save_config(cfg)
    return jsonify({'success': True})


@bp.route('/models')
def list_models():
    cfg = _load_config()
    models_dir = cfg.get('nllb_dir', '')
    models = []
    if models_dir and os.path.isdir(models_dir):
        for name in sorted(os.listdir(models_dir)):
            full = os.path.join(models_dir, name)
            if os.path.isdir(full) and os.path.exists(os.path.join(full, 'config.json')):
                models.append({'name': name, 'type': 'huggingface', 'path': full})
    return jsonify({'models': models, 'models_dir': models_dir, 'last_model': cfg.get('last_nllb_model', '')})


@bp.route('/languages')
def list_languages():
    langs = [{'code': c, 'name': n} for c, n in LANGUAGES.items() if c != 'auto']
    return jsonify({'src_langs': [{'code': 'auto', 'name': '自动检测'}] + langs, 'tgt_langs': langs})


@bp.route('/unload-model', methods=['POST'])
def unload_model():
    with _model_lock:
        _unload_model()
    cfg = _load_config()
    cfg['last_nllb_model'] = ''
    _save_config(cfg)
    return jsonify({'success': True, 'message': '模型已卸载'})


@bp.route('/load-model', methods=['POST'])
def load_model():
    data = request.get_json() or {}
    model_name = data.get('model_name', '').strip()
    if not model_name:
        return jsonify({'success': False, 'error': '请指定模型名称'}), 400

    cfg = _load_config()
    models_dir = cfg.get('nllb_dir', '')
    if not models_dir or not os.path.isdir(models_dir):
        return jsonify({'success': False, 'error': '请先配置 NLLB 模型目录'}), 400

    model_path = os.path.join(models_dir, model_name)
    if not os.path.isdir(model_path):
        return jsonify({'success': False, 'error': f'模型目录不存在: {model_name}'}), 400

    with _model_lock:
        _unload_model()
        ok, err = _load_model(model_path, model_name)
        if ok:
            cfg['last_nllb_model'] = model_name
            _save_config(cfg)
            return jsonify({'success': True, 'model_name': model_name})
        return jsonify({'success': False, 'error': err or '加载失败'}), 400


@bp.route('/translate-text', methods=['POST'])
def translate_text():
    data = request.get_json()
    text = (data.get('text') or '').strip()
    if not text:
        return jsonify({'success': False, 'error': '请输入要翻译的文本'}), 400

    src_lang = data.get('src_lang', 'en')
    tgt_lang = data.get('tgt_lang', 'zh')
    cfg = _load_config()
    batch_size = cfg.get('batch_size', 8)

    cfg['last_src_lang'] = src_lang
    cfg['last_tgt_lang'] = tgt_lang
    _save_config(cfg)

    ok, err = _ensure_model(src_lang, tgt_lang)
    if not ok:
        return jsonify({'success': False, 'error': err}), 400

    task_id = uuid.uuid4().hex
    _tasks[task_id] = {'status': 'translating', 'progress': 0, 'result': None, 'error': None}

    def _do_task():
        try:
            paragraphs = [p.strip() for p in text.split('\n') if p.strip()]
            if not paragraphs:
                paragraphs = [text]

            all_results = []
            total = len(paragraphs)
            for i in range(0, total, batch_size):
                batch = paragraphs[i:i + batch_size]
                results, err = _do_translate(batch, src_lang, tgt_lang, cfg)
                if err:
                    _tasks[task_id] = {'status': 'error', 'error': err}
                    return
                all_results.extend(results)
                _tasks[task_id]['progress'] = round((i + len(batch)) / total * 100)

            _tasks[task_id] = {'status': 'done', 'result': '\n'.join(all_results), 'progress': 100}

        except Exception as e:
            _tasks[task_id] = {'status': 'error', 'error': str(e)}

    threading.Thread(target=_do_task, daemon=True).start()
    return jsonify({'success': True, 'task_id': task_id})


@bp.route('/translate-file', methods=['POST'])
def translate_file():
    if 'subtitle' not in request.files:
        return jsonify({'success': False, 'error': '请上传字幕文件'}), 400
    f = request.files['subtitle']
    if not f.filename:
        return jsonify({'success': False, 'error': '请上传字幕文件'}), 400

    ext = os.path.splitext(f.filename)[1].lower()
    if ext not in SUBTITLE_EXTS:
        return jsonify({'success': False, 'error': f'不支持的格式: {ext}，支持 SRT/ASS/VTT'}), 400

    src_lang = request.form.get('src_lang', 'en')
    tgt_lang = request.form.get('tgt_lang', 'zh')
    cfg = _load_config()
    batch_size = cfg.get('batch_size', 8)

    cfg['last_src_lang'] = src_lang
    cfg['last_tgt_lang'] = tgt_lang
    _save_config(cfg)

    ok, err = _ensure_model(src_lang, tgt_lang)
    if not ok:
        return jsonify({'success': False, 'error': err}), 400

    content = f.read().decode('utf-8', errors='replace')
    fmt = _detect_subtitle_format(f.filename, content)

    task_id = uuid.uuid4().hex
    _tasks[task_id] = {'status': 'translating', 'progress': 0, 'result': None, 'error': None,
                       'filename': f.filename, 'format': fmt}

    def _do_task():
        try:
            if fmt == 'srt':
                entries = _parse_srt(content)
                texts = [e[2] for e in entries]
            elif fmt == 'vtt':
                entries = _parse_vtt(content)
                texts = [e[2] for e in entries]
            elif fmt == 'ass':
                full_lines = content.split('\n')
                dialogue_lines = [l for l in full_lines if l.strip().lower().startswith('dialogue:')]
                texts = []
                for dl in dialogue_lines:
                    _, txt = _ass_extract_text(dl)
                    texts.append(txt if txt else '')
                entries = dialogue_lines
            else:
                _tasks[task_id] = {'status': 'error', 'error': '未知字幕格式'}
                return

            if not texts:
                _tasks[task_id] = {'status': 'error', 'error': '字幕文件中未找到可翻译的文本'}
                return

            all_results = []
            total = len(texts)

            for i in range(0, total, batch_size):
                batch = texts[i:i + batch_size]
                batch = [t for t in batch if t.strip()]
                if not batch:
                    all_results.extend(texts[i:i + batch_size])
                    _tasks[task_id]['progress'] = round((i + batch_size) / total * 100)
                    continue

                results, err = _do_translate(batch, src_lang, tgt_lang, cfg)
                if err:
                    _tasks[task_id] = {'status': 'error', 'error': err}
                    return

                ri = 0
                for j in range(i, min(i + batch_size, total)):
                    if texts[j].strip():
                        all_results.append(results[ri] if ri < len(results) else texts[j])
                        ri += 1
                    else:
                        all_results.append('')
                _tasks[task_id]['progress'] = round(min(i + batch_size, total) / total * 100)

            if fmt == 'srt':
                translated_entries = [(entries[i][0], entries[i][1], all_results[i]) for i in range(len(entries))]
                output = _build_srt(translated_entries)
            elif fmt == 'vtt':
                translated_entries = [(entries[i][0], entries[i][1], all_results[i]) for i in range(len(entries))]
                output = _build_vtt(translated_entries)
            elif fmt == 'ass':
                output = _build_ass(content.split('\n'), all_results)
            else:
                output = '\n'.join(all_results)

            out_path = os.path.join(tempfile.gettempdir(), f'ai_trans_{task_id}{ext}')
            with open(out_path, 'w', encoding='utf-8') as fout:
                fout.write(output)

            _tasks[task_id] = {
                'status': 'done', 'progress': 100,
                'result': output, 'file_path': out_path,
                'filename': f.filename, 'format': fmt,
            }

        except Exception as e:
            _tasks[task_id] = {'status': 'error', 'error': str(e)}

    threading.Thread(target=_do_task, daemon=True).start()
    return jsonify({'success': True, 'task_id': task_id})


@bp.route('/task-status/<task_id>')
def task_status(task_id):
    task = _tasks.get(task_id)
    if not task:
        return jsonify({'success': False, 'error': '任务不存在'}), 404
    return jsonify({'success': True, **task})


@bp.route('/download/<task_id>')
def download(task_id):
    task = _tasks.get(task_id)
    if not task or task.get('status') != 'done':
        return jsonify({'error': '文件不存在'}), 404
    file_path = task.get('file_path')
    if not file_path or not os.path.exists(file_path):
        return jsonify({'error': '文件不存在'}), 404
    orig = task.get('filename', 'translated')
    base, ext = os.path.splitext(orig)
    return send_file(file_path, as_attachment=True, download_name=f'{base}_translated{ext}')


@bp.route('/cleanup/<task_id>', methods=['POST'])
def cleanup(task_id):
    task = _tasks.pop(task_id, None)
    if task and task.get('file_path'):
        try:
            os.remove(task['file_path'])
        except OSError:
            pass
    return jsonify({'success': True})