#!/usr/bin/env python3
"""
L5X RLL Text-Format Parser

Parses the compact text rung format from Allen-Bradley .L5X exports into
structured data (Networks → Rungs → Branches → Elements).

Text format grammar:
  Rung ::= (Instruction | Branch)* ';'
  Branch ::= '[' BranchPath (',' BranchPath)* ']'
  BranchPath ::= Instruction*
  Instruction ::= OTE | OTL | OTU | XIC | XIO | TON | TOF | RTO | CTU | CTD |
                 RES | MOV | COP | BTD | LIM | CPT | EQU | NEQ | GRT | LES |
                 GEQ | LEQ | ADD | SUB | MUL | DIV | NOP | MCR | JSR | SRT |
                 SST | CLR | OTS | FLL | SWPB | FBC | GSV | Z_* | Technifor_EIP
                 | AnyUnknown3Char

Contacts (inline logic): XIC, XIO
Outputs (right rail): OTE, OTL, OTU
Function blocks: everything else
"""

import re
import xml.etree.ElementTree as ET
from typing import List, Dict, Any, Optional


# ── Element classification ───────────────────────────────────────────

CONTACTS = {'XIC', 'XIO'}
COILS = {'OTE', 'OTL', 'OTU', 'OTS'}
TIMER_TYPES = {'TON', 'TOF', 'RTO'}
COUNTER_TYPES = {'CTU', 'CTD', 'CTUD'}
COMPARISONS = {'EQU', 'NEQ', 'GRT', 'LES', 'GEQ', 'LEQ', 'LIM'}
DATA_OPS = {'MOV', 'COP', 'BTD', 'CPT', 'CLR', 'FLL', 'SWPB', 'EXG', 'FBC', 'GSV',
            'SQT', 'ABS', 'NEG', 'INC', 'DEC'}
MATH_OPS = {'ADD', 'SUB', 'MUL', 'DIV'}
CONTROL_OPS = {'NOP', 'MCR', 'JSR', 'SRT', 'SST', 'RES', 'JMP', 'LBL', 'ONS', 'OSR', 'OSF'}

ALL_TYPES = CONTACTS | COILS | TIMER_TYPES | COUNTER_TYPES | COMPARISONS | DATA_OPS | MATH_OPS | CONTROL_OPS


# ── Tokenizer ────────────────────────────────────────────────────────

def tokenize_rung(text: str) -> List[str]:
    """
    Split rung text into individual instruction tokens.
    Handles: INST(args), INST(args)INST2(args2), [branch1,branch2], nested [...]
    """
    tokens = []
    i = 0
    text = text.strip().rstrip(';').strip()
    
    while i < len(text):
        if text[i] == ' ':
            i += 1
            continue
        
        if text[i] == '[':
            # Find matching closing bracket
            depth = 0
            start = i
            while i < len(text):
                if text[i] == '[':
                    depth += 1
                elif text[i] == ']':
                    depth -= 1
                    if depth == 0:
                        break
                i += 1
            tokens.append(text[start:i+1])
            i += 1
            continue
        
        if text[i] == ',':
            tokens.append(',')
            i += 1
            continue
        
        # Match instruction: NAME(args) or NAME();
        match = re.match(r'([A-Za-z_][A-Za-z0-9_]*)(?:\(([^)]*)\))?', text[i:])
        if match:
            name = match.group(1)
            args_str = match.group(2)
            if name in ALL_TYPES or (len(name) <= 3 and name.isupper()):
                tokens.append(f"{name}({args_str})" if args_str is not None else name)
                i += match.end()
                continue
        
        # Skip unknown characters
        i += 1
    
    return tokens


# ── Instruction Parser ───────────────────────────────────────────────

def parse_instruction(token: str) -> Dict[str, Any]:
    """
    Parse a single instruction token into a structured element dict.
    
    Returns:
        {
            'type': 'XIC',
            'comment': 'tag_name',
            'tag': 'tag_name',
            'value': value,        # for comparisons
            'params': [args],      # raw parameter list
            'is_contact': bool,
            'is_coil': bool,
            'is_block': bool
        }
    """
    # Extract name and args
    match = re.match(r'([A-Za-z_][A-Za-z0-9_]*)(?:\((.*)\))?', token)
    if not match:
        return {'type': 'UNKNOWN', 'comment': token, 'params': [], 'is_block': True}
    
    name = match.group(1)
    args_str = match.group(2) or ''
    params = [p.strip() for p in args_str.split(',') if p.strip()] if args_str else []
    
    result = {
        'type': name,
        'params': params,
    }
    
    if name in CONTACTS:
        result['is_contact'] = True
        result['is_coil'] = False
        result['is_block'] = False
        result['tag'] = params[0] if params else ''
        result['comment'] = params[0] if params else ''
    
    elif name in COILS:
        result['is_contact'] = False
        result['is_coil'] = True
        result['is_block'] = False
        result['tag'] = params[0] if params else ''
        result['comment'] = params[0] if params else ''
    
    elif name in TIMER_TYPES:
        result['is_contact'] = False
        result['is_coil'] = False
        result['is_block'] = True
        result['tag'] = params[0] if params else ''
        result['comment'] = params[0] if params else ''
        # TON(timer, preset, base) — preset may be '?' in L5X
        if len(params) >= 2:
            try:
                result['preset'] = int(params[1]) if params[1] != '?' else 0
            except ValueError:
                result['preset'] = 0
    
    elif name in COUNTER_TYPES:
        result['is_contact'] = False
        result['is_coil'] = False
        result['is_block'] = True
        result['tag'] = params[0] if params else ''
        result['comment'] = params[0] if params else ''
        if len(params) >= 2:
            try:
                result['preset'] = int(params[1]) if params[1] != '?' else 0
            except ValueError:
                result['preset'] = 0
    
    elif name in COMPARISONS:
        result['is_contact'] = True  # comparisons act as contacts
        result['is_coil'] = False
        result['is_block'] = False
        result['tag'] = params[0] if len(params) >= 1 else ''
        if len(params) >= 2:
            try:
                result['value'] = int(params[1])
            except ValueError:
                try:
                    result['value'] = float(params[1])
                except ValueError:
                    result['value'] = params[1]
        if len(params) >= 3:
            try:
                result['value2'] = int(params[2])
            except ValueError:
                try:
                    result['value2'] = float(params[2])
                except ValueError:
                    result['value2'] = params[2]
        result['comment'] = name
    
    elif name in DATA_OPS | MATH_OPS | CONTROL_OPS:
        result['is_contact'] = False
        result['is_coil'] = False
        result['is_block'] = True
        result['tag'] = params[0] if params else ''
        result['comment'] = name
        # Store additional params as value/value2
        if len(params) >= 2:
            try:
                result['value'] = int(params[1])
            except ValueError:
                result['value'] = params[1]
        if len(params) >= 3:
            try:
                result['value2'] = int(params[2])
            except ValueError:
                result['value2'] = params[2]
    
    else:
        # Unknown instruction — treat as function block
        result['is_contact'] = False
        result['is_coil'] = False
        result['is_block'] = True
        result['tag'] = params[0] if params else ''
        result['comment'] = name
    
    return result


# ── Branch Parser ────────────────────────────────────────────────────

def parse_branch_group(text: str) -> List[List[Dict[str, Any]]]:
    """
    Parse a bracket-enclosed branch group: [path1 ,path2 ,path3]
    
    Returns list of branches, each branch is a list of elements.
    """
    # Strip outer brackets
    inner = text.strip()[1:-1].strip()
    
    # Split by comma at top level (not inside brackets)
    paths = []
    depth = 0
    current = []
    for ch in inner:
        if ch == '[':
            depth += 1
            current.append(ch)
        elif ch == ']':
            depth -= 1
            current.append(ch)
        elif ch == ',' and depth == 0:
            paths.append(''.join(current))
            current = []
        else:
            current.append(ch)
    
    if current:
        paths.append(''.join(current))
    
    branches = []
    for path in paths:
        path = path.strip()
        if not path:
            continue
        
        # Tokenize and parse this path
        tokens = tokenize_rung(path)
        elements = []
        for tok in tokens:
            if tok == ',':
                continue
            if tok.startswith('['):
                # Nested branch
                nested = parse_branch_group(tok)
                # Flatten nested for simplicity — add as a sub-branch marker
                for sub_branch in nested:
                    elements.extend(sub_branch)
            else:
                elem = parse_instruction(tok)
                elements.append(elem)
        branches.append(elements)
    
    return branches


# ── Rung Parser ──────────────────────────────────────────────────────

def parse_rung(text: str) -> Dict[str, Any]:
    """
    Parse a complete rung text into structured data.
    
    Returns:
        {
            'branches': [[elements], [elements], ...],  # parallel branches
            'serial': [elements],                        # elements after branches
            'raw': original_text
        }
    """
    text = text.strip().rstrip(';').strip()
    if not text or text == 'NOP':
        return {'branches': [], 'serial': [], 'raw': text}
    
    tokens = tokenize_rung(text)
    
    # Group tokens: before-branches (serial), branches, after-branches (serial)
    before_serial = []
    branch_groups = []
    after_serial = []
    
    in_branches = False
    current_branch_group = []
    
    i = 0
    while i < len(tokens):
        tok = tokens[i]
        
        if tok.startswith('['):
            in_branches = True
            # Find the complete branch group token
            current_branch_group = [tok]
            branch_groups.append(parse_branch_group(tok))
            i += 1
        elif tok == ',' and in_branches:
            # Comma between branches at top level — handled by parse_branch_group
            i += 1
        elif tok.startswith(']') and in_branches:
            # End of branch group
            in_branches = False
            i += 1
        else:
            elem = parse_instruction(tok)
            if in_branches:
                # Elements after closing bracket but before next branch
                after_serial.append(elem)
            else:
                before_serial.append(elem)
            i += 1
    
    # Combine: if there are branches, the before_serial are lead-in contacts
    # and after_serial are trailing elements (coils, blocks)
    # Flatten branch groups: each group is [[path1_elements], [path2_elements]]
    # We want a flat list of all branch paths across all groups
    flat_branches = []
    for group in branch_groups:
        flat_branches.extend(group)

    result = {
        'lead_in': before_serial,
        'branches': flat_branches,
        'trailing': after_serial,
        'raw': text
    }
    
    return result


# ── L5X File Parser ──────────────────────────────────────────────────

def parse_l5x_programs(l5x_path: str) -> Dict[str, Any]:
    """
    Parse an L5X file and extract the full Program → Routine → Rung hierarchy.
    
    Returns:
        {
            'programs': [
                {
                    'name': program_name,
                    'routines': [
                        {
                            'name': routine_name,
                            'type': 'RLL' | 'LAD',
                            'rung_count': int,
                            'rungs': [parsed_rung_dict, ...]
                        }
                    ],
                    'routine_count': int
                }
            ],
            'program_count': int
        }
    """
    tree = ET.parse(l5x_path)
    root = tree.getroot()
    
    programs = []
    
    # Navigate: Controller → Programs → Program
    controller = root.find('Controller')
    if controller is None:
        return {'programs': [], 'program_count': 0}
    
    programs_el = controller.find('Programs')
    if programs_el is None:
        return {'programs': [], 'program_count': 0}
    
    for program in programs_el:
        program_name = program.get('Name')
        if not program_name:
            name_el = program.find('Name')
            program_name = name_el.text if name_el is not None and name_el.text else 'Unknown'
        
        # Find Routines within this Program
        routines_el = program.find('Routines')
        if routines_el is None:
            continue
        
        routine_list = []
        for routine in routines_el:
            routine_name = routine.get('Name')
            if not routine_name:
                name_el = routine.find('Name')
                routine_name = name_el.text if name_el is not None and name_el.text else 'Unknown'
            
            lang = routine.get('Type')
            if not lang:
                lang_el = routine.find('Language')
                lang = lang_el.text if lang_el is not None and lang_el.text else 'Unknown'
            
            # Accept both 'RLL' and 'LAD' as ladder logic types
            if lang not in ('RLL', 'LAD'):
                continue
            
            # For RLL, rungs are in RLLContent
            rung_list = []
            content_el = routine.find('RLLContent')
            if content_el is not None:
                for rung in content_el:
                    if rung.tag != 'Rung':
                        continue
                    
                    num_str = rung.get('Number')
                    if not num_str:
                        num_el = rung.find('Number')
                        num_str = num_el.text if num_el is not None and num_el.text else '0'
                    
                    comment_el = rung.find('Comment')
                    comment = comment_el.text if comment_el is not None and comment_el.text else ''
                    
                    text_el = rung.find('Text')
                    text = text_el.text if text_el is not None else ''
                    text = text.strip().rstrip(';').strip()
                    
                    num = int(num_str) if num_str else 0
                    parsed = parse_rung(text)
                    
                    rung_list.append({
                        'number': num,
                        'comment': comment,
                        'parsed': parsed,
                        'raw_text': text
                    })
            
            rung_list.sort(key=lambda r: r['number'])
            
            routine_list.append({
                'name': routine_name,
                'type': lang,
                'rung_count': len(rung_list),
                'rungs': rung_list,
            })
        
        if routine_list:
            programs.append({
                'name': program_name,
                'routines': routine_list,
                'routine_count': len(routine_list),
            })
    
    return {
        'programs': programs,
        'program_count': len(programs),
    }


def parse_l5x_routines(l5x_path: str) -> List[Dict[str, Any]]:
    """
    Parse an L5X file and extract all RLL routines with their rungs.
    
    Legacy flat list — prefer parse_l5x_programs() for full hierarchy.
    """
    result = parse_l5x_programs(l5x_path)
    routines = []
    for prog in result['programs']:
        for rout in prog['routines']:
            routines.append({
                'program': prog['name'],
                'name': rout['name'],
                'language': rout['type'],
                'rung_count': rout['rung_count'],
                'rungs': rout['rungs'],
            })
    return routines


def parse_l5x_to_program(l5x_path: str, routine_index: int = 0) -> Dict[str, Any]:
    """
    Parse L5X and return the first (or specified) routine as a program dict
    compatible with the existing ladder interpreter format.
    """
    routines = parse_l5x_routines(l5x_path)
    if not routines:
        return {'name': 'Empty', 'networks': []}
    
    routine = routines[routine_index]
    
    # Convert to Networks format for the ladder interpreter
    networks = []
    for rung in routine['rungs']:
        parsed = rung['parsed']
        
        # Build elements list
        elements = list(parsed.get('lead_in', []))
        
        # If there are branches, represent as parallel paths
        branches = parsed.get('branches', [])
        trailing = parsed.get('trailing', [])
        
        network = {
            'number': rung['number'],
            'comment': rung['comment'],
            'elements': elements,
            'parallel': branches,
            'outputs': trailing,
            'raw': rung['raw_text']
        }
        networks.append(network)
    
    return {
        'name': routine['name'],
        'networks': networks,
        'routine_count': len(routines),
        'all_routines': [r['name'] for r in routines],
    }