import tkinter as tk
from tkinter import ttk, filedialog, scrolledtext
import struct
import os

class FileHeaderParser:
    def __init__(self, root):
        self.root = root
        self.root.title("文件头解析器")
        self.root.geometry("1000x600")  # 增加窗口宽度
        
        # 加载模板
        self.templates = {}
        self.load_templates()
        
        # 文件扩展名到模板名称的映射
        self.extension_to_template = {
            # 图像格式
            '.bmp': 'BMP文件头',
            '.png': 'PNG文件头',
            '.jpg': 'JPG文件头',
            '.jpeg': 'JPG文件头',
            '.gif': 'GIF文件头',
            '.tiff': 'TIFF文件头',
            '.tif': 'TIFF文件头',
            
            # 音频格式
            '.mp3': 'MP3文件头',
            '.wav': 'WAV文件头',
            
            # 视频格式
            '.mp4': 'MP4文件头',
            '.avi': 'AVI文件头',
            
            # 文档格式
            '.pdf': 'PDF文件头',
            '.doc': 'DOC文件头',
            '.docx': 'DOC文件头',
            '.xls': 'XLS文件头',
            '.xlsx': 'XLS文件头',
            '.ppt': 'PPT文件头',
            '.pptx': 'PPT文件头',
            
            # 压缩格式
            '.zip': 'ZIP文件头',
            
            # 可执行文件格式
            '.exe': 'PE文件头',
            '.dll': 'PE文件头',
            '.sys': 'PE文件头',
            '.drv': 'PE文件头',
            '.ocx': 'PE文件头',
            '.cpl': 'PE文件头',
            '.scr': 'PE文件头'
        }
        
        # 顶部区域
        self.top_frame = ttk.Frame(root, padding="5")
        self.top_frame.pack(fill=tk.X)
        
        # 文件选择
        self.file_frame = ttk.Frame(self.top_frame)
        self.file_frame.pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        ttk.Label(self.file_frame, text="文件:").pack(side=tk.LEFT, padx=(0, 5))
        
        self.file_path = tk.StringVar()
        self.file_entry = ttk.Entry(self.file_frame, textvariable=self.file_path)
        self.file_entry.pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        self.browse_button = ttk.Button(self.file_frame, text="选择文件", command=self.browse_file)
        self.browse_button.pack(side=tk.LEFT, padx=5)
        
        # 解析按钮
        self.parse_button = ttk.Button(self.top_frame, text="解析文件头", command=self.parse_header)
        self.parse_button.pack(side=tk.RIGHT, padx=5)
        
        # 中间区域（左右布局）
        self.middle_frame = ttk.Frame(root, padding="5")
        self.middle_frame.pack(fill=tk.BOTH, expand=True)
        
        # 左侧区域
        self.left_frame = ttk.Frame(self.middle_frame, width=300)  # 固定左侧宽度
        self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH)
        self.left_frame.pack_propagate(False)  # 防止自动调整大小
        
        # 模板选择
        self.template_frame = ttk.Frame(self.left_frame)
        self.template_frame.pack(fill=tk.X, pady=5)
        
        ttk.Label(self.template_frame, text="选择模板:").pack(side=tk.LEFT)
        self.template_var = tk.StringVar()
        self.template_combo = ttk.Combobox(self.template_frame, textvariable=self.template_var, values=list(self.templates.keys()))
        self.template_combo.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        self.template_combo.bind('<<ComboboxSelected>>', self.on_template_selected)
        
        # 模板内容
        ttk.Label(self.left_frame, text="模板内容:").pack(anchor=tk.W, pady=5)
        self.format_text = scrolledtext.ScrolledText(self.left_frame, width=35)  # 减小文本框宽度
        self.format_text.pack(fill=tk.BOTH, expand=True)
        
        # 右侧输出
        self.right_frame = ttk.Frame(self.middle_frame)
        self.right_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=(10, 0))
        
        ttk.Label(self.right_frame, text="解析结果:").pack(anchor=tk.W)
        
        # 创建Treeview控件
        self.result_tree = ttk.Treeview(self.right_frame, columns=('field', 'type', 'value', 'hex'), show='headings')
        
        # 设置列标题
        self.result_tree.heading('field', text='字段名')
        self.result_tree.heading('type', text='类型[字节数]')
        self.result_tree.heading('value', text='值')
        self.result_tree.heading('hex', text='原始数据(HEX)')
        
        # 设置列宽
        self.result_tree.column('field', width=150, minwidth=100)
        self.result_tree.column('type', width=100, minwidth=80)
        self.result_tree.column('value', width=150, minwidth=100)
        self.result_tree.column('hex', width=200, minwidth=150)
        
        # 添加滚动条
        scrollbar = ttk.Scrollbar(self.right_frame, orient=tk.VERTICAL, command=self.result_tree.yview)
        self.result_tree.configure(yscrollcommand=scrollbar.set)
        
        # 放置控件
        self.result_tree.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        
        # 绑定列宽调整事件
        self.result_tree.bind('<Configure>', self.on_tree_configure)
        
    def load_templates(self):
        try:
            with open('templates.txt', 'r', encoding='utf-8') as f:
                content = f.read()
                current_template = None
                current_content = []
                
                for line in content.split('\n'):
                    line = line.strip()
                    if not line:
                        if current_template and current_content:
                            self.templates[current_template] = '\n'.join(current_content)
                            current_content = []
                        continue
                        
                    if line.startswith('[') and line.endswith(']'):
                        if current_template and current_content:
                            self.templates[current_template] = '\n'.join(current_content)
                        current_template = line[1:-1]
                        current_content = []
                    else:
                        current_content.append(line)
                
                if current_template and current_content:
                    self.templates[current_template] = '\n'.join(current_content)
        except Exception as e:
            print(f"加载模板文件失败: {str(e)}")
            self.templates = {"错误": "加载模板失败"}
            
    def on_tree_configure(self, event):
        # 获取Treeview的宽度
        tree_width = self.result_tree.winfo_width()
        
        # 计算每列的宽度比例
        total_width = tree_width - 20  # 减去滚动条宽度
        self.result_tree.column('field', width=int(total_width * 0.25))
        self.result_tree.column('type', width=int(total_width * 0.15))
        self.result_tree.column('value', width=int(total_width * 0.25))
        self.result_tree.column('hex', width=int(total_width * 0.35))
        
    def on_template_selected(self, event):
        template_name = self.template_var.get()
        if template_name in self.templates:
            self.format_text.delete(1.0, tk.END)
            self.format_text.insert(1.0, self.templates[template_name])
            
    def select_template_by_extension(self, file_path):
        ext = os.path.splitext(file_path)[1].lower()
        if ext in self.extension_to_template:
            template_name = self.extension_to_template[ext]
            self.template_var.set(template_name)
            self.on_template_selected(None)
        
    def browse_file(self):
        filename = filedialog.askopenfilename()
        if filename:
            self.file_path.set(filename)
            self.select_template_by_extension(filename)
            
    def parse_header(self):
        try:
            # 获取文件路径
            file_path = self.file_path.get()
            if not file_path:
                self.result_tree.delete(*self.result_tree.get_children())
                return
                
            # 获取解析格式
            format_text = self.format_text.get(1.0, tk.END).strip()
            if not format_text:
                self.result_tree.delete(*self.result_tree.get_children())
                return
                
            # 计算需要读取的总字节数
            total_bytes = 0
            format_list = []
            for line in format_text.split('\n'):
                if not line.strip():
                    continue
                parts = line.strip().split(',')
                type_name = parts[0].strip()
                size = int(parts[1].strip())
                field_name = parts[2].strip() if len(parts) > 2 else ""
                total_bytes += size
                format_list.append((type_name, size, field_name))
                
            # 读取文件头
            with open(file_path, 'rb') as f:
                header_data = f.read(total_bytes)
                
            # 解析并显示结果
            self.result_tree.delete(*self.result_tree.get_children())
            offset = 0
            for type_name, size, field_name in format_list:
                data = header_data[offset:offset+size]
                hex_data = ''.join([f'{b:02X}' for b in data])
                
                if type_name == 'byte':
                    result = ' '.join([f'{b:02X}' for b in data])
                elif type_name == 'num':
                    if size == 4:
                        result = str(struct.unpack('<I', data)[0])  # 小端模式
                    elif size == 2:
                        result = str(struct.unpack('<H', data)[0])  # 小端模式
                    elif size == 1:
                        result = str(data[0])
                    else:
                        result = f"不支持的数值大小: {size}"
                elif type_name == 'char':
                    result = data.decode('ascii', errors='replace')
                else:
                    result = f"不支持的类型: {type_name}"
                    
                field_display = field_name if field_name else ""
                self.result_tree.insert('', 'end', values=(field_display, f"{type_name}[{size}]", result, hex_data))
                offset += size
                
        except Exception as e:
            self.result_tree.delete(*self.result_tree.get_children())
            self.result_tree.insert('', 'end', values=('错误', '', str(e), ''))

if __name__ == "__main__":
    root = tk.Tk()
    app = FileHeaderParser(root)
    root.mainloop()