#!/usr/bin/env python3
"""
Convert ANTLR4 g4 grammar to Lark format.

This converter handles ANTLR4 grammar files and produces Lark grammar.
"""

import re
from pathlib import Path
from dataclasses import dataclass
from typing import Optional


@dataclass
class Rule:
    name: str
    body: str
    is_lexer: bool
    is_fragment: bool = False
    channel: Optional[str] = None


class ANTLR4ToLark:
    """Convert ANTLR4 grammar to Lark format."""

    def __init__(self):
        self.rules: dict[str, Rule] = {}
        self.options: dict[str, str] = {}
        self.token_vocab: Optional[str] = None

    def parse_g4(self, content: str, filename: str = "") -> list[Rule]:
        """Parse an ANTLR4 grammar file."""
        rules = []

        # Remove comments
        content = re.sub(r'//[^\n]*', '', content)
        content = re.sub(r'/\*[\s\S]*?\*/', '', content)

        # Extract grammar type and name
        grammar_match = re.search(r'(lexer\s+|parser\s+)?grammar\s+(\w+)\s*;', content)
        is_lexer_grammar = grammar_match and 'lexer' in (grammar_match.group(1) or '')

        # Extract options
        options_match = re.search(r'options\s*\{([^}]*)\}', content)
        if options_match:
            opts = options_match.group(1)
            for opt in re.findall(r'(\w+)\s*=\s*(\w+)', opts):
                self.options[opt[0]] = opt[1]
            if 'tokenVocab' in self.options:
                self.token_vocab = self.options['tokenVocab']

        # Find all rules - need to handle quoted strings properly
        # Process line by line to properly handle rule boundaries
        lines = content.split('\n')
        current_rule = None
        current_body = []
        pending_name = None  # For rules where name is on separate line from :

        for line in lines:
            stripped = line.strip()

            # Skip empty lines and comments
            if not stripped or stripped.startswith('//'):
                continue

            # Check for rule name on its own line (ANTLR4 format style)
            # e.g., "script" followed by "    : query SEMI? EOF"
            name_only = re.match(r'^(fragment\s+)?([A-Za-z_]\w*)\s*$', stripped)
            if name_only and not current_rule:
                pending_name = (name_only.group(2), bool(name_only.group(1)))
                continue

            # Check for : at start of line (continuation of pending rule)
            if stripped.startswith(':') and pending_name:
                name, is_frag = pending_name
                rest = stripped[1:].strip()  # Remove leading :
                pending_name = None

                if self._ends_with_semicolon(rest):
                    body = rest.rstrip(';').strip()
                    self._add_rule((name, body, is_frag), is_lexer_grammar, rules)
                else:
                    current_rule = (name, '', is_frag)
                    current_body = [rest]
                continue

            # Check for inline rule start (name : body on same line)
            rule_start = re.match(r'^(fragment\s+)?([A-Za-z_]\w*)\s*:\s*(.*)', stripped)
            if rule_start and not current_rule:
                pending_name = None  # Clear any pending
                is_frag = bool(rule_start.group(1))
                name = rule_start.group(2)
                rest = rule_start.group(3)

                if self._ends_with_semicolon(rest):
                    body = rest.rstrip(';').strip()
                    self._add_rule((name, body, is_frag), is_lexer_grammar, rules)
                else:
                    current_rule = (name, '', is_frag)
                    current_body = [rest]
                continue

            # Continue building current rule body
            if current_rule:
                # Handle | at start of line (alternative)
                if stripped.startswith('|'):
                    current_body.append(stripped)
                else:
                    current_body.append(stripped)

                full_body = ' '.join(current_body)
                if self._ends_with_semicolon(full_body):
                    body = full_body.rstrip(';').strip()
                    name, _, is_frag = current_rule
                    self._add_rule((name, body, is_frag), is_lexer_grammar, rules)
                    current_rule = None
                    current_body = []

        return rules

    def _ends_with_semicolon(self, text: str) -> bool:
        """Check if text ends with semicolon outside of quotes."""
        in_single = False
        in_double = False
        for i, c in enumerate(text):
            if c == "'" and not in_double:
                in_single = not in_single
            elif c == '"' and not in_single:
                in_double = not in_double
        # Check if ends with ; and not in quotes
        return text.rstrip().endswith(';') and not in_single and not in_double

    def _add_rule(self, rule_tuple, is_lexer_grammar, rules):
        """Add a parsed rule to the rules list."""
        name, body, is_fragment = rule_tuple
        is_lexer = name[0].isupper() or is_lexer_grammar

        # Check for channel directive
        channel = None
        channel_match = re.search(r'->\s*channel\s*\(\s*(\w+)\s*\)', body)
        if channel_match:
            channel = channel_match.group(1)
            body = re.sub(r'\s*->\s*channel\s*\([^)]+\)', '', body)

        # Check for skip directive
        if re.search(r'->\s*skip', body, re.IGNORECASE):
            body = re.sub(r'\s*->\s*skip', '', body, flags=re.IGNORECASE)
            channel = 'skip'

        rule = Rule(
            name=name,
            body=body.strip(),
            is_lexer=is_lexer,
            is_fragment=is_fragment,
            channel=channel,
        )
        rules.append(rule)
        self.rules[name] = rule

    def to_snake_case(self, name: str) -> str:
        """Convert camelCase or PascalCase to snake_case."""
        # Remove oC_ prefix if present
        if name.startswith('oC_'):
            name = name[3:]
        # Insert underscore before uppercase letters
        result = re.sub(r'(?<!^)(?=[A-Z])', '_', name)
        return result.lower()

    def convert_body(self, body: str, is_lexer: bool) -> str:
        """Convert ANTLR4 rule body to Lark syntax."""
        result = body

        # Handle alternatives spanning multiple lines
        result = re.sub(r'\s+\|\s+', ' | ', result)
        result = re.sub(r'\s+', ' ', result)

        # FIRST: Convert string literals to Lark format (before any other transformations)
        # This preserves the content of string literals
        if self.options.get('caseInsensitive') == 'true':
            # Convert all string literals to case-insensitive
            def make_case_insensitive(m):
                s = m.group(1) if m.group(1) is not None else m.group(2)
                # Escape any special characters
                return f'"{s}"i'
            result = re.sub(r"'([^']*)'|\"([^\"]*)\"", make_case_insensitive, result)
        else:
            # Just convert single quotes to double quotes
            def convert_quotes(m):
                s = m.group(1)
                return f'"{s}"'
            result = re.sub(r"'([^']*)'", convert_quotes, result)

        # Convert ANTLR4 range operator .. to Lark regex (for character ranges)
        # ANTLR4: 'a'..'z' or "a".."z"
        # Lark: /[a-z]/
        result = re.sub(r'"([^"]+)"\s*\.\.\s*"([^"]+)"', r'/[\1-\2]/', result)

        # Convert ~ (negation in ANTLR4 lexer) to regex negation
        result = re.sub(r'~\s*\[([^\]]+)\]', r'/[^\1]/', result)
        result = re.sub(r'~\s*"([^"]*)"', r'/[^\1]/', result)

        # Convert . (any char) to Lark regex - but NOT inside strings or regexes
        # Only match standalone dots, not those in already-converted strings
        result = re.sub(r'(?<!["\'/\[\w])\s*\.\s*(?!["\'/\]\w*+?])', ' /./ ', result)

        # Convert rule references to snake_case (for parser rules)
        def convert_rule_ref(m):
            name = m.group(1)
            # Skip if inside a string or regex
            if m.group(0) != name:
                return m.group(0)
            # Check if it's a known rule
            if name in self.rules:
                rule = self.rules[name]
                if rule.is_lexer or rule.is_fragment:
                    # Lexer rule - keep uppercase
                    return name.upper() if not name.isupper() else name
                else:
                    # Parser rule - convert to snake_case
                    return self.to_snake_case(name)
            elif name.isupper():
                # Unknown uppercase - assume lexer terminal
                return name
            elif any(c.isupper() for c in name):
                # Has uppercase - convert to snake_case
                return self.to_snake_case(name)
            else:
                # Already lowercase
                return name

        # Find rule references (not in strings)
        result = re.sub(r'\b([a-zA-Z_][a-zA-Z0-9_]*)\b(?=(?:[^"]*"[^"]*")*[^"]*$)', convert_rule_ref, result)

        # Convert non-greedy quantifiers (Lark handles greediness differently)
        result = re.sub(r'\*\?', '*', result)
        result = re.sub(r'\+\?', '+', result)
        result = re.sub(r'\?\?', '?', result)

        # Handle semantic predicates {condition}? - Lark doesn't support
        result = re.sub(r'\{[^}]+\}\s*\?', '', result)

        # Handle actions {code} - remove for grammar conversion
        result = re.sub(r'\{[^}]+\}', '', result)

        # Clean up whitespace
        result = re.sub(r'\s+', ' ', result).strip()

        return result

    def to_lark(self) -> str:
        """Generate complete Lark grammar."""
        lines = [
            "// Cypher Grammar - Auto-generated from ANTLR4",
            "// Source: grammars-v4/cypher",
            "",
        ]

        # Separate parser and lexer rules
        parser_rules = []
        lexer_rules = []
        fragments = []
        skip_rules = []

        for name, rule in self.rules.items():
            if rule.channel == 'skip':
                skip_rules.append((name, rule))
            elif rule.is_fragment:
                fragments.append((name, rule))
            elif rule.is_lexer:
                lexer_rules.append((name, rule))
            else:
                parser_rules.append((name, rule))

        # Find start rule
        start_candidates = ['script', 'cypher', 'query', 'statement', 'oC_Cypher']
        start_rule = None
        for candidate in start_candidates:
            if candidate in self.rules:
                start_rule = self.to_snake_case(candidate)
                break

        if start_rule:
            lines.append(f"start: {start_rule}")
            lines.append("")

        # Parser rules
        lines.append("// === Parser Rules ===")
        lines.append("")

        for name, rule in parser_rules:
            rule_name = self.to_snake_case(name)
            body = self.convert_body(rule.body, is_lexer=False)
            lines.append(f"{rule_name}: {body}")

        lines.append("")
        lines.append("// === Lexer Rules (Terminals) ===")
        lines.append("")

        for name, rule in lexer_rules:
            body = self.convert_body(rule.body, is_lexer=True)
            lines.append(f"{name}: {body}")

        if fragments:
            lines.append("")
            lines.append("// === Fragments ===")
            lines.append("")

            for name, rule in fragments:
                body = self.convert_body(rule.body, is_lexer=True)
                # In Lark, fragments are just regular terminals
                lines.append(f"{name}: {body}")

        # Add ignore directives for skip rules
        if skip_rules:
            lines.append("")
            lines.append("// === Ignored Tokens ===")
            for name, rule in skip_rules:
                body = self.convert_body(rule.body, is_lexer=True)
                lines.append(f"{name}: {body}")
                lines.append(f"%ignore {name}")

        return '\n'.join(lines)

    def convert_files(self, input_paths: list[Path], output_path: Path):
        """Convert multiple g4 files and write combined output."""
        for path in input_paths:
            content = path.read_text()
            self.parse_g4(content, path.name)

        lark_content = self.to_lark()
        output_path.write_text(lark_content)

        print(f"Converted {len(self.rules)} rules from {len(input_paths)} files")
        print(f"Output: {output_path}")

        # Summary
        parser_count = sum(1 for r in self.rules.values() if not r.is_lexer and not r.is_fragment)
        lexer_count = sum(1 for r in self.rules.values() if r.is_lexer and not r.is_fragment)
        fragment_count = sum(1 for r in self.rules.values() if r.is_fragment)

        print(f"\nRule types:")
        print(f"  Parser rules: {parser_count}")
        print(f"  Lexer rules: {lexer_count}")
        print(f"  Fragments: {fragment_count}")


def main():
    antlr_dir = Path(__file__).parent / 'antlr4'
    output_dir = Path(__file__).parent / 'lark_grammars'

    # Convert grammars-v4 version (split lexer/parser)
    print("=== Converting grammars-v4 (CypherLexer + CypherParser) ===")
    converter1 = ANTLR4ToLark()
    converter1.convert_files(
        [antlr_dir / 'CypherLexer.g4', antlr_dir / 'CypherParser.g4'],
        output_dir / 'cypher_from_grammars_v4.lark'
    )

    print()

    # Convert OpenCypher M14 version (combined grammar)
    print("=== Converting OpenCypher M14 (Cypher.g4) ===")
    converter2 = ANTLR4ToLark()
    converter2.convert_files(
        [antlr_dir / 'Cypher_OpenCypher.g4'],
        output_dir / 'cypher_from_opencypher.lark'
    )


if __name__ == '__main__':
    main()
