import os
import re
import csv
import argparse
from datetime import datetime

'''
使用示例：
python3 scan_blocks.py \
  --project_root "/Users/**/Documents/ProjectDemo" \
  --output_dir "/Users/**/Downloads" \
  --prefixes DD \
  --suffixes .m .mm
'''

# 忽略關鍵字
ignore_block_keywords = ['mas_', 'snp.', 'UIView animate', 'UIView.animate', "enumerate", "async", "TZI"]

# 注釋判斷
def disable_line(line):
    dis_line = r'^\s*\/\/'
    return re.match(dis_line, line) is not None

# 提取 block 內容
def extract_block(lines, start_idx):
    block_lines = []
    open_brace_count = 0
    in_block = False
    end_idx = 0
    for i in range(start_idx, len(lines)):
        line = lines[i]
        block_lines.append(line)
        open_brace_count += line.count('{')
        open_brace_count -= line.count('}')
        if '{' in line:
            in_block = True
        if in_block and open_brace_count == 0:
            end_idx = i
            break
    return (''.join(block_lines), end_idx)

# 檢測 block 是否安全,安全則過濾
def block_is_safe(block_str):
    # 找到第一個 '^' 的位置
    block_index = block_str.find('^')
    block_content = block_str[block_index:-1]
    # 找到^之后第一個 '{' 的位置
    a = block_content.find('{')
    if a == -1:
        # 沒有 '{'
        return True
    # 找到最后一個 '}' 的位置
    b = block_content.rfind('}')
    if b == -1 or a >= b:
        # 沒有 '}'
        return True
    # 提取 { 和 } 之間的內容
    substring = block_content[a + 1:b]

    self_positions = []
    start = 0
    while True:
        pos = substring.lower().find('self', start)
        if pos == -1:
            break
        self_positions.append(pos)
        start = pos + 1

    if len(self_positions) == 0:
        # 不存在self
        return True
    # 檢查每個self
    all_strong_weak = True
    for pos in self_positions:
        if pos >= 1 and substring[pos - 1:pos] in [' ', ':', '[', ',']:
            all_strong_weak = False
        elif pos >= 1 and substring[pos - 1:pos] == '(':
            if pos >= 7 and substring[pos - 7:pos].lower() == 'typeof(':
                pass
            else:
                all_strong_weak = False
        elif pos < 1:
            all_strong_weak = False

    if all_strong_weak:
        return True
    # 過濾沒匹配到 strongify
    save_pattern = r'(strongify|zf_strongify)'
    return bool(re.findall(save_pattern, substring))

# 檢查是否為方法聲明
def check_method_declaration(lines, start_idx):
    method_declaration_start = r'^\s*[\-+]'
    # 初始行是否是方法聲明，不是則直接返回
    if re.match(method_declaration_start, lines[start_idx]) is None:
        return (None, start_idx)
    method_declaration_pattern = r'^\s*[\-+]\s*\(.*\)\s*.*\{'
    method_lines = []
    end_idx = -1
    for i in range(start_idx, len(lines)):
        method_lines.append(lines[i].strip())
        if '{' in lines[i]:
            end_idx = i
            break
    method_declaration = ' '.join(method_lines)
    return (re.match(method_declaration_pattern, method_declaration), end_idx)

# 掃描單個文件
def scan_file_for_blocks(file_path):
    try:
        with open(file_path, "r", encoding='utf-8', errors='ignore') as f:
            lines = f.readlines()
    except Exception as e:
        print(f"[ERROR] Failed to read {file_path}: {e}")
        return []

    suspicious_blocks = []
    line_start_index = -1
    for idx, line in enumerate(lines):
        if line_start_index >= idx:
            continue
        if disable_line(line):
            continue
        if any(keyword.lower() in line.lower() for keyword in ignore_block_keywords):
            continue
        # 如果是方法聲明，忽略該行
        (declar_result, end_idx) = check_method_declaration(lines, idx)
        if declar_result is not None:
            line_start_index = end_idx
            continue
        # 檢測是否是忽略關鍵字
        if '^' in line and "@property" not in line and "typedef" not in line:
            (block_content, end_idx) = extract_block(lines, idx)
            line_start_index = end_idx
            if 'self' in block_content and not block_is_safe(block_content):
                suspicious_blocks.append((idx + 1, block_content.strip()))
    return suspicious_blocks

# 掃描整個項目目錄
def scan_project(project_root, deal_dir_Path, ignore_dir_Path, deal_file_prefix, deal_file_subfix):
    results = []
    for root, dirs, files in os.walk(project_root):
        # 如果存在忽略的文件，過濾忽略的文件
        if len(ignore_dir_Path) and any(root.startswith(ig) for ig in ignore_dir_Path):
            continue
        # 如果存在要處理的文件，只處理要處理的文件
        if len(deal_dir_Path) and not any(root.startswith(dp) for dp in deal_dir_Path):
            continue
        for file in files:
            deal_file = False
            if any(file.endswith(ext) for ext in deal_file_subfix):
                if deal_file_prefix:
                    if any(file.lower().startswith(pre.lower()) for pre in deal_file_prefix):
                        deal_file = True
                else:
                    deal_file = True
            if deal_file:
                file_path = os.path.join(root, file)
                result = scan_file_for_blocks(file_path)
                if result:
                    print("??" * 50)
                    print(f"\n[!]{file_path}:")
                    for line_num, content in result:
                        results.append([file_path, line_num, content])
                        print(f"  Line {line_num}:\n{content}\n")
    return results

# 保存結果為 CSV
def save_results_to_csv(results, output_dir, csv_file='block'):
    os.makedirs(output_dir, exist_ok=True)
    full_path = os.path.join(output_dir, csv_file + datetime.now().strftime("%Y%m%d%H%M%S") + ".csv")
    # 寫入CSV文件
    with open(full_path, mode='w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(["File Path", "Line Number", "Block Content"])
        writer.writerows(results)
    print(f"Results saved to {full_path}")

# 命令行參數(shù)
def parse_args():
    parser = argparse.ArgumentParser(description="Scan Objective-C project for unsafe blocks.")
    parser.add_argument('--project_root', type=str, required=True, help='項目路徑')
    parser.add_argument('--output_dir', type=str, default='./output', help='輸出文件路徑')
    parser.add_argument('--deal_dirs', type=str, nargs='*', default=[], help='處理所有pods文件夾下的內容,比如處理所有Pods/中的文件,為空則默認全部都處理')
    parser.add_argument('--ignore_dirs', type=str, nargs='*', default=[], help='忽略處理所有pods文件夾下的內容.為空則默認全部不忽略，不為空則只忽略指定的')
    parser.add_argument('--prefixes', type=str, nargs='*', default=[], help='需要處理的文件前綴 例如所有DD開頭的文件：DD，為空則所有文件')
    parser.add_argument('--suffixes', type=str, nargs='*', default=['.m'], help='需要處理的文件后綴 例如所有.m后綴的文件：.m')
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    results = scan_project(
        project_root=args.project_root,
        deal_dir_Path=args.deal_dirs,
        ignore_dir_Path=args.ignore_dirs,
        deal_file_prefix=args.prefixes,
        deal_file_subfix=args.suffixes
    )
    if results:
        save_results_to_csv(results, args.output_dir)
    else:
        print("Empty Results.")