#!/usr/bin/env python3
"""
Batch Translation Script for Chinese Comments

Translates Chinese comments in code files to English using LLM API.

Usage:
    # Translate a single file
    python scripts/translate_chinese_comments.py --file path/to/file.py
    
    # Translate all files in a directory
    python scripts/translate_chinese_comments.py --dir chengyue_script/dvlm
    
    # Dry run (show changes without modifying files)
    python scripts/translate_chinese_comments.py --dir chengyue_script --dry-run
    
    # List files with Chinese content
    python scripts/translate_chinese_comments.py --dir chengyue_script --list-only
"""

import argparse
import os
import re
import sys
from pathlib import Path
from typing import List, Tuple, Optional

# Chinese character pattern
CHINESE_PATTERN = re.compile(r'[\u4e00-\u9fff]+')

# Translation mappings for common phrases (to avoid API calls for simple translations)
COMMON_TRANSLATIONS = {
    # Training related
    "训练": "training",
    "模型": "model",
    "参数": "parameter",
    "数据": "data",
    "数据集": "dataset",
    "批量": "batch",
    "学习率": "learning rate",
    "损失": "loss",
    "梯度": "gradient",
    "检查点": "checkpoint",
    "恢复": "resume",
    "保存": "save",
    "加载": "load",
    "配置": "configuration",
    "环境": "environment",
    "变量": "variable",
    "输出": "output",
    "输入": "input",
    "日志": "log",
    "调试": "debug",
    "测试": "test",
    "评测": "evaluation",
    "推理": "inference",
    
    # MDM/VLM specific
    "掩码": "mask",
    "块": "block",
    "块大小": "block size",
    "扩散": "diffusion",
    "自回归": "autoregressive",
    "视觉": "vision",
    "语言": "language",
    "多模态": "multimodal",
    "图像": "image",
    "文本": "text",
    "处理器": "processor",
    "分词器": "tokenizer",
    
    # Common phrases
    "开始": "start",
    "结束": "end",
    "成功": "success",
    "失败": "failed",
    "错误": "error",
    "警告": "warning",
    "信息": "info",
    "找到": "found",
    "没有找到": "not found",
    "存在": "exists",
    "不存在": "does not exist",
    "注意": "note",
    "重要": "important",
    "可选": "optional",
    "必需": "required",
    "默认": "default",
    "自动": "auto",
    "手动": "manual",
    "启用": "enable",
    "禁用": "disable",
    "激活": "activate",
    "最新": "latest",
    "版本": "version",
    
    # Slurm/HPC
    "节点": "node",
    "多节点": "multi-node",
    "单节点": "single-node",
    "显存": "GPU memory",
    "内存": "memory",
    "任务": "task",
    "作业": "job",
    "提交": "submit",
    "队列": "queue",
    "超时": "timeout",
    "信号": "signal",
    
    # Comments
    "修改点": "modification point",
    "适用": "applicable for",
    "推荐": "recommended",
    "备选": "alternative",
    "评测工具": "evaluation tool",
    "工具": "tool",
    "说明": "description",
    "类型": "type",
    "样本数": "sample count",
    "评测方式": "evaluation method",
    "选择题": "multiple choice",
    "数学推理": "math reasoning",
}


def has_chinese(text: str) -> bool:
    """Check if text contains Chinese characters."""
    return bool(CHINESE_PATTERN.search(text))


def translate_simple(text: str) -> str:
    """Translate using common mappings (no API call)."""
    result = text
    for chinese, english in sorted(COMMON_TRANSLATIONS.items(), key=lambda x: -len(x[0])):
        result = result.replace(chinese, english)
    return result


def translate_with_llm(text: str, provider: str = "anthropic") -> str:
    """Translate using LLM API."""
    try:
        from cursor_agent.tools.llm_api import call_llm
        
        prompt = f"""Translate the following Chinese text to English. Keep any code, variable names, or technical terms unchanged. Only translate the natural language parts.

Text to translate:
{text}

Translated text:"""
        
        result = call_llm(prompt, provider=provider)
        return result.strip()
    except ImportError:
        # Fallback to simple translation
        return translate_simple(text)
    except Exception as e:
        print(f"LLM translation failed: {e}", file=sys.stderr)
        return translate_simple(text)


def find_chinese_in_line(line: str) -> List[Tuple[int, int, str]]:
    """Find all Chinese text segments in a line.
    
    Returns list of (start, end, text) tuples.
    """
    matches = []
    for match in CHINESE_PATTERN.finditer(line):
        # Expand to include surrounding Chinese characters and punctuation
        start = match.start()
        end = match.end()
        
        # Expand backwards
        while start > 0 and (has_chinese(line[start-1]) or line[start-1] in '，。：；""''【】（）'):
            start -= 1
        
        # Expand forwards  
        while end < len(line) and (has_chinese(line[end]) or line[end] in '，。：；""''【】（）'):
            end += 1
            
        # Avoid duplicates
        if not any(s <= start < e or s < end <= e for s, e, _ in matches):
            matches.append((start, end, line[start:end]))
    
    return matches


def translate_line(line: str, use_llm: bool = False) -> str:
    """Translate Chinese text in a single line."""
    if not has_chinese(line):
        return line
    
    matches = find_chinese_in_line(line)
    if not matches:
        return line
    
    result = line
    # Process matches in reverse order to preserve indices
    for start, end, chinese_text in sorted(matches, reverse=True):
        if use_llm:
            translated = translate_with_llm(chinese_text)
        else:
            translated = translate_simple(chinese_text)
        
        # Only replace if translation is different and complete
        if translated != chinese_text and not has_chinese(translated):
            result = result[:start] + translated + result[end:]
    
    return result


def translate_file(filepath: Path, dry_run: bool = False, use_llm: bool = False) -> Tuple[bool, int]:
    """Translate Chinese comments in a file.
    
    Returns (modified, count) where count is number of lines changed.
    """
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except UnicodeDecodeError:
        print(f"Skipping binary file: {filepath}", file=sys.stderr)
        return False, 0
    except Exception as e:
        print(f"Error reading {filepath}: {e}", file=sys.stderr)
        return False, 0
    
    modified_lines = []
    change_count = 0
    
    for i, line in enumerate(lines):
        if has_chinese(line):
            new_line = translate_line(line, use_llm=use_llm)
            if new_line != line:
                change_count += 1
                if dry_run:
                    print(f"{filepath}:{i+1}")
                    print(f"  - {line.rstrip()}")
                    print(f"  + {new_line.rstrip()}")
            modified_lines.append(new_line)
        else:
            modified_lines.append(line)
    
    if change_count > 0 and not dry_run:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.writelines(modified_lines)
    
    return change_count > 0, change_count


def find_files_with_chinese(directory: Path, extensions: List[str] = None) -> List[Path]:
    """Find all files with Chinese content in a directory."""
    if extensions is None:
        extensions = ['.py', '.sh', '.md', '.txt', '.json']
    
    files = []
    for ext in extensions:
        for filepath in directory.rglob(f'*{ext}'):
            if filepath.is_file() and not any(p.startswith('.') for p in filepath.parts):
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        content = f.read()
                    if has_chinese(content):
                        files.append(filepath)
                except:
                    pass
    
    return sorted(files)


def main():
    parser = argparse.ArgumentParser(
        description="Translate Chinese comments in code files to English",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    parser.add_argument('--file', type=Path, help='Single file to translate')
    parser.add_argument('--dir', type=Path, help='Directory to scan for files')
    parser.add_argument('--dry-run', action='store_true', help='Show changes without modifying files')
    parser.add_argument('--list-only', action='store_true', help='Only list files with Chinese content')
    parser.add_argument('--use-llm', action='store_true', help='Use LLM API for translation (slower but more accurate)')
    parser.add_argument('--extensions', type=str, default='.py,.sh,.md,.txt', 
                       help='Comma-separated file extensions to process')
    
    args = parser.parse_args()
    
    if not args.file and not args.dir:
        parser.error("Either --file or --dir must be specified")
    
    extensions = [ext.strip() if ext.startswith('.') else f'.{ext.strip()}' 
                  for ext in args.extensions.split(',')]
    
    if args.file:
        files = [args.file]
    else:
        files = find_files_with_chinese(args.dir, extensions)
    
    if args.list_only:
        print(f"Found {len(files)} files with Chinese content:")
        for f in files:
            print(f"  {f}")
        return
    
    total_files = 0
    total_lines = 0
    
    for filepath in files:
        modified, count = translate_file(filepath, dry_run=args.dry_run, use_llm=args.use_llm)
        if modified:
            total_files += 1
            total_lines += count
            if not args.dry_run:
                print(f"Translated {count} lines in {filepath}")
    
    print(f"\n{'Would modify' if args.dry_run else 'Modified'} {total_lines} lines in {total_files} files")


if __name__ == '__main__':
    main()

