#!/usr/bin/env python3
# ASR1806 Partition Configuration Generator (Final Correct Format)

import re
import json
import sys
import os
import argparse
from collections import OrderedDict
from pathlib import Path

def validate_file(path, description):
    """Validate input file exists and is readable"""
    path_obj = Path(path)
    if not path_obj.exists():
        raise FileNotFoundError(f"{description} not found: {path}")
    if not path_obj.is_file():
        raise IOError(f"{description} path is not a file: {path}")
    if not os.access(str(path_obj), os.R_OK):
        raise PermissionError(f"Cannot read {description}: {path}")
    return path_obj

def parse_dtsi_partitions(dtsi_file):
    """Parse partition information from DTSI file"""
    partitions = OrderedDict()
    try:
        with open(dtsi_file, 'r', encoding='utf-8') as f:
            content = f.read()
        
        pattern = re.compile(
            r'partition@([0-9A-Fa-f]+)\s*{\s*.*?label\s*=\s*"([^"]+)";\s*.*?reg\s*=\s*<([^>]+)>;\s*.*?};',
            re.DOTALL
        )
        
        for match in pattern.finditer(content):
            try:
                start_addr = int(match.group(1), 16)
                label = match.group(2).strip()
                reg_values = [int(x, 16) for x in match.group(3).split()]
                
                if not reg_values:
                    print(f"Invalid reg value for partition {label}", file=sys.stderr)
                    continue
                    
                size = reg_values[-1] if len(reg_values) > 1 else reg_values[0]
                
                partitions[label] = {
                    'start_addr': f"0x{start_addr:08X}",
                    'size': f"0x{size:08X}",
                    'size_bytes': size
                }
            except ValueError as e:
                print(f"Failed to parse partition data: {e}", file=sys.stderr)
                continue
        
        if not partitions:
            raise ValueError("No valid partitions found in DTSI file")
            
        print(f"Parsed {len(partitions)} partitions from {dtsi_file}")
        return partitions
    except Exception as e:
        raise RuntimeError(f"DTSI parse failed: {str(e)}") from e

def load_config(config_file):
    """Load and validate partition configuration"""
    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f, object_pairs_hook=OrderedDict)
        
        required = {
            'mtdparts_order': list,
            'partition_aliases': dict
        }
        for key, typ in required.items():
            if key not in config:
                raise KeyError(f"Missing required key: {key}")
            if not isinstance(config[key], typ):
                raise TypeError(f"{key} must be {typ.__name__}")
        
        print(f"Loaded config from {config_file}")
        return config
    except json.JSONDecodeError as e:
        raise RuntimeError(f"Invalid JSON format: {str(e)}") from e
    except Exception as e:
        raise RuntimeError(f"Config load failed: {str(e)}") from e

def generate_complete_config(partitions, config):
    """Generate the complete configuration block with exact formatting"""
    # 1. Generate partition defines
    defines = []
    for part_name, alias in config['partition_aliases'].items():
        if part_name not in partitions:
            print(f"Config references unknown partition: {part_name}", file=sys.stderr)
            continue
        defines.append(f'#define {alias} "{part_name}"')
    defines_section = "\n".join(defines)

    # 2. Generate MTDPARTS content with exact formatting
    mtdparts = []
    valid_parts = [p for p in config['mtdparts_order'] if p in partitions and p in config['partition_aliases']]
    num_parts = len(valid_parts)
    
    for i, part_name in enumerate(valid_parts):
        part = partitions[part_name]
        alias = config['partition_aliases'][part_name]
        
        # 构建每个分区的定义，格式为 "0x00180000@0x02600000("UBOOT_MTD_PART_A"),"
        mtd_entry = f'0x{part["size"][2:]}@0x{part["start_addr"][2:]}("{alias}")'
        if i < num_parts - 1:
            mtd_entry += ','  # 除最后一个分区外，每个分区后加逗号
        
        # 每个分区定义用双引号包裹，并添加行继续符和缩进
        if i == 0:
            # 第一个分区不需要额外缩进
            mtdparts.append(f'"mtdparts=nand_mtd:{mtd_entry}"\\')
        else:
            # 后续分区需要 5 个空格缩进
            mtdparts.append(f'     "{mtd_entry}"\\')
    
    # 最后一行需要单独处理，确保没有双引号和反斜杠
    if mtdparts:
        last_line = mtdparts.pop()  # 移除最后一行
        # 构建不带双引号和反斜杠的最后一行
        mtdparts.append(last_line.replace('\\"', '').replace('\\', ''))
    
    # 连接所有分区定义
    mtdparts_joined = '\n'.join(mtdparts)
    
    # 构建最终的 MTDPARTS 部分
    mtdparts_section = f"""#ifdef CONFIG_CMD_FASTBOOT
#define CONFIG_MTDPARTS \\
{mtdparts_joined}
#endif /* CONFIG_CMD_FASTBOOT */"""

    # 3. Combine into complete block
    complete_content = f"""{defines_section}

{mtdparts_section}"""
    
    return complete_content

def update_header_file(header_file, complete_content):
    """Replace ONLY the content between dynamic markers"""
    try:
        header_file.parent.mkdir(parents=True, exist_ok=True)
        
        if not header_file.exists():
            raise RuntimeError(f"Header file does not exist: {header_file}")
        
        # Read existing content
        with open(header_file, 'r', encoding='utf-8') as f:
            content = f.read()
        original_mode = header_file.stat().st_mode
        
        # Pattern to find the exact block to replace
        pattern = re.compile(
            r'(/\* add mtd parts dynamically \*/\n)(.*?)(\n#else)',
            re.DOTALL
        )
        
        match = pattern.search(content)
        if not match:
            raise RuntimeError("Could not find the replacement markers ('/* add mtd parts dynamically */' to '#else')")
        
        # Replace only the middle part
        updated_content = (
            content[:match.start(2)] + 
            complete_content + 
            content[match.end(2):]
        )
        
        # Write updated content
        with open(header_file, 'w', encoding='utf-8') as f:
            f.write(updated_content)
        header_file.chmod(original_mode)
        
        print(f"Successfully updated {header_file}")
        return True
        
    except Exception as e:
        raise RuntimeError(f"Header update failed: {str(e)}") from e

def main():
    parser = argparse.ArgumentParser(
        description='Generate partition configuration between markers',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    
    parser.add_argument('--dtsi', required=True,
                       help='Input DTSI file with partition layout')
    parser.add_argument('--config', required=True,
                       help='JSON config file with partition aliases')
    parser.add_argument('--header', required=True,
                       help='Output header file path')
    parser.add_argument('--verbose', action='store_true',
                       help='Show detailed processing info')
    
    args = parser.parse_args()
    
    try:
        dtsi_path = validate_file(args.dtsi, "DTSI file")
        config_path = validate_file(args.config, "Config file")
        header_path = Path(args.header).resolve()
        
        if args.verbose:
            print(f"Processing:\n- DTSI: {dtsi_path}\n- Config: {config_path}\n- Output: {header_path}")
        
        partitions = parse_dtsi_partitions(dtsi_path)
        config = load_config(config_path)
        
        # Verify all configured partitions exist
        missing = [p for p in config['mtdparts_order'] if p not in partitions]
        if missing:
            print(f"Missing partitions in DTSI: {', '.join(missing)}", file=sys.stderr)
        
        complete_content = generate_complete_config(partitions, config)
        
        if args.verbose:
            print("\nGenerated content to insert:")
            print("-" * 40)
            print(complete_content)
            print("-" * 40)
        
        updated = update_header_file(header_path, complete_content)
        
        if updated:
            print("Update completed successfully")
        
        return 0
        
    except Exception as e:
        print(f"Error: {str(e)}", file=sys.stderr)
        return 1

if __name__ == "__main__":
    sys.exit(main())