#!/usr/bin/env python3
"""
Generate Rust test files from Wasot syntax test markdown files.

This script parses markdown files containing test cases and generates
standard Rust test files that can be run with `cargo test`.
"""

import re
import sys
import subprocess
import argparse
from pathlib import Path
from typing import List, Dict, Any

from parse_syntax_file import TestCase, MarkdownTestParser


class RustTestGenerator:
    """Generates Rust test code from test cases."""
    
    def __init__(self, test_cases: List[TestCase]):
        self.test_cases = test_cases
    
    def generate(self) -> str:
        """Generate complete Rust test file."""
        output = []
        
        # File header
        output.append(self._generate_header())
        
        # Group tests by chapter
        chapters = {}
        for test in self.test_cases:
            chapter = test.chapter or "uncategorized"
            if chapter not in chapters:
                chapters[chapter] = []
            chapters[chapter].append(test)
        
        # Generate test modules for each chapter
        for chapter_name, tests in chapters.items():
            output.append(self._generate_chapter_module(chapter_name, tests))
        
        # Close the main module
        output.append('}')
        
        return '\n'.join(output)
    
    def _generate_header(self) -> str:
        """Generate file header with imports."""
        return '''// This file is auto-generated by scripts/generate_syntax_tests.py
// DO NOT EDIT MANUALLY

#[cfg(test)]
#[allow(dead_code, unused_imports, unused)]
mod generated_syntax_tests {
    use wasot_compiler::events::{Event, Emitter};
    use wasot_compiler::parsers::wasot::{Footnote, Link, MetaLine, QuotedString, Id};
    use wasot_compiler::emitters::wasot_file_emitter::WasotFileEmitter;
    use wasot_compiler::parsers::wasot::arguments::ParsedArguments;
    use std::sync::mpsc;
    use std::collections::HashMap;
'''
    
    def _generate_chapter_module(self, chapter_name: str, tests: List[TestCase]) -> str:
        """Generate a test module for a chapter."""
        # Convert chapter name to valid Rust module name
        module_name = re.sub(r'[^a-zA-Z0-9_]', '_', chapter_name.lower())
        module_name = re.sub(r'_+', '_', module_name).strip('_')
        
        output = [f"\n    mod {module_name} {{"]
        output.append("        use super::*;")
        
        for test in tests:
            output.append(self._generate_test_function(test))
        
        output.append("    }")
        return '\n'.join(output)
    
    def _generate_test_function(self, test: TestCase) -> str:
        """Generate a single test function."""
        # Convert test name to valid Rust function name
        fn_name = re.sub(r'[^a-zA-Z0-9_]', '_', test.name.lower())
        fn_name = re.sub(r'_+', '_', fn_name).strip('_')
        
        # Make sure name doesn't start with number
        if fn_name and fn_name[0].isdigit():
            fn_name = 'test_' + fn_name
        
        output = [f"\n        #[test]"]
        output.append(f"        fn {fn_name}() {{")
        
        # Add description as comment
        if test.description:
            for line in test.description.split('\n'):
                output.append(f"            // {line}")
        
        # Generate input
        output.append(f"            let input = {self._rust_string_literal(test.input)};")
        output.append("")
        
        # Parse and collect events
        output.append("            // Parse input and collect events")
        output.append("            let emitter = WasotFileEmitter::from_string(input.to_string());")
        output.append("            let receiver = emitter.subscribe();")
        output.append("            ")
        output.append("            std::thread::spawn(move || {")
        output.append("                emitter.run();")
        output.append("            });")
        output.append("            ")
        output.append("            let mut events = Vec::new();")
        output.append("            for event in receiver {")
        output.append("                events.push(event);")
        output.append("            }")
        output.append("")
        
        # Generate assertions
        output.append("            // Verify events")
        output.append(f"            assert_eq!(events.len(), {len(test.expected_events)}, \"Expected {len(test.expected_events)} events, got {{}}\", events.len());")
        
        for i, expected in enumerate(test.expected_events):
            output.append(f"            ")
            output.append(f"            // Event {i}")
            output.extend(self._generate_event_assertion(i, expected))
        
        output.append("        }")
        return '\n'.join(output)
    
    def _generate_event_assertion(self, index: int, expected: Any) -> List[str]:
        """Generate assertion for a single event."""
        output = []
        
        if isinstance(expected, dict):
            # Single event type in dict
            event_type, event_data = next(iter(expected.items()))
            
            if event_type == "block-started":
                output.extend(self._generate_block_started_assertion(index, event_data))
            elif event_type == "void-element":
                output.extend(self._generate_void_element_assertion(index, event_data))
            elif event_type == "text-finished":
                output.extend(self._generate_text_finished_assertion(index, event_data))
            elif event_type == "block-finished":
                output.append(f"            assert!(matches!(&events[{index}], Event::BlockFinished));")
            elif event_type == "text":
                # Shorthand for text-finished with just lines
                output.extend(self._generate_text_finished_assertion(index, {"lines": event_data if isinstance(event_data, list) else [event_data]}))
            elif event_type == "block":
                # Shorthand for block-started with just type
                output.extend(self._generate_block_started_assertion(index, event_data))
            elif event_type == "void":
                # Shorthand for void-element with just type
                output.extend(self._generate_void_element_assertion(index, event_data))
        elif isinstance(expected, str):
            # Simple string event (e.g., "block-finished")
            if expected == "block-finished":
                output.append(f"            assert!(matches!(&events[{index}], Event::BlockFinished));")
            elif expected.startswith("block:"):
                # Shorthand like "block:div" 
                block_type = expected[6:]
                output.extend(self._generate_block_started_assertion(index, block_type))
        
        return output
    
    def _generate_block_started_assertion(self, index: int, data: Dict[str, Any]) -> List[str]:
        """Generate assertions for BlockStarted event.
        
        Only asserts on values explicitly provided in the YAML.
        This makes tests more focused on what they're actually testing.
        """
        output = []
        
        # Handle shorthand string format: just the type name
        if isinstance(data, str):
            output.append(f"            match &events[{index}] {{")
            output.append("                Event::BlockStarted { meta_line } => {")
            output.append(f"                    assert_eq!(meta_line.type_name, {self._rust_string_literal(data)});")
            output.append("                }")
            output.append(f"                _ => panic!(\"Expected BlockStarted at index {index}, got {{:?}}\", events[{index}]),")
            output.append("            }")
            return output
        
        output.append(f"            match &events[{index}] {{")
        output.append("                Event::BlockStarted { meta_line } => {")
        
        # Only assert on explicitly provided values
        if 'type' in data:
            type_name = data['type']
            output.append(f"                    assert_eq!(meta_line.type_name, {self._rust_string_literal(type_name)});")
        
        if 'level' in data:
            level = data['level']
            output.append(f"                    assert_eq!(meta_line.level, {level});")
        
        if 'id' in data:
            id_value = data['id']
            if id_value:
                output.append(f"                    assert_eq!(meta_line.id.value, {self._rust_string_literal(id_value)});")
            else:
                output.append("                    assert!(meta_line.id.is_empty());")
        
        if 'inline-text' in data:
            inline_text = data['inline-text']
            if inline_text:
                output.append(f"                    assert_eq!(meta_line.inline_text.content, {self._rust_string_literal(inline_text)});")
            else:
                output.append("                    assert!(meta_line.inline_text.content.is_empty());")

        if 'metaline-quoted' in data:
            quoted = data['metaline-quoted']
            if quoted:
                output.append(f"                    assert_eq!(meta_line.inline_text.quoted, true);")
            else:
                output.append(f"                    assert_eq!(meta_line.inline_text.quoted, false);")

        if 'type-quoted' in data:
            type_quoted = data['type-quoted']
            if type_quoted:
                output.append(f"                    assert_eq!(meta_line.type_quoted, true);")
            else:
                output.append(f"                    assert_eq!(meta_line.type_quoted, false);")

        if 'args' in data:
            args = data['args']
            if args:
                output.append(f"                    assert_eq!(meta_line.arguments.args.len(), {len(args)});")
                for i, arg in enumerate(args):
                    output.append(f"                    assert_eq!(meta_line.arguments.args[{i}].content, {self._rust_string_literal(arg)});")
            else:
                output.append("                    assert_eq!(meta_line.arguments.args.len(), 0);")
        
        if 'kwargs' in data:
            kwargs = data['kwargs']
            if kwargs:
                output.append(f"                    assert_eq!(meta_line.arguments.kwargs.len(), {len(kwargs)});")
                for key, value in kwargs.items():
                    output.append(f"                    assert_eq!(meta_line.arguments.kwargs.iter().find(|(k, _)| k.content == {self._rust_string_literal(key)}).unwrap().1.content, {self._rust_string_literal(value)});")
            else:
                output.append("                    assert_eq!(meta_line.arguments.kwargs.len(), 0);")
        
        output.append("                }")
        output.append(f"                _ => panic!(\"Expected BlockStarted at index {index}, got {{:?}}\", events[{index}]),")
        output.append("            }")

        return output

    def _generate_void_element_assertion(self, index: int, data: Dict[str, Any]) -> List[str]:
        """Generate assertions for VoidElement event.

        Only asserts on values explicitly provided in the YAML.
        This is very similar to BlockStarted but matches on VoidElement event type.
        """
        output = []

        # Handle shorthand string format: just the type name
        if isinstance(data, str):
            output.append(f"            match &events[{index}] {{")
            output.append("                Event::VoidElement { meta_line } => {")
            output.append(f"                    assert_eq!(meta_line.type_name, {self._rust_string_literal(data)});")
            output.append("                }")
            output.append(f"                _ => panic!(\"Expected VoidElement at index {index}, got {{:?}}\", events[{index}]),")
            output.append("            }")
            return output

        output.append(f"            match &events[{index}] {{")
        output.append("                Event::VoidElement { meta_line } => {")

        # Only assert on explicitly provided values
        if 'type' in data:
            type_name = data['type']
            output.append(f"                    assert_eq!(meta_line.type_name, {self._rust_string_literal(type_name)});")

        if 'level' in data:
            level = data['level']
            output.append(f"                    assert_eq!(meta_line.level, {level});")

        if 'id' in data:
            id_value = data['id']
            if id_value:
                output.append(f"                    assert_eq!(meta_line.id.value, {self._rust_string_literal(id_value)});")
            else:
                output.append("                    assert!(meta_line.id.is_empty());")

        if 'inline-text' in data:
            inline_text = data['inline-text']
            if inline_text:
                output.append(f"                    assert_eq!(meta_line.inline_text.content, {self._rust_string_literal(inline_text)});")
            else:
                output.append("                    assert!(meta_line.inline_text.content.is_empty());")

        if 'metaline-quoted' in data:
            quoted = data['metaline-quoted']
            if quoted:
                output.append(f"                    assert_eq!(meta_line.inline_text.quoted, true);")
            else:
                output.append(f"                    assert_eq!(meta_line.inline_text.quoted, false);")

        if 'type-quoted' in data:
            type_quoted = data['type-quoted']
            if type_quoted:
                output.append(f"                    assert_eq!(meta_line.type_quoted, true);")
            else:
                output.append(f"                    assert_eq!(meta_line.type_quoted, false);")

        if 'args' in data:
            args = data['args']
            if args:
                output.append(f"                    assert_eq!(meta_line.arguments.args.len(), {len(args)});")
                for i, arg in enumerate(args):
                    output.append(f"                    assert_eq!(meta_line.arguments.args[{i}].content, {self._rust_string_literal(arg)});")
            else:
                output.append("                    assert_eq!(meta_line.arguments.args.len(), 0);")

        if 'kwargs' in data:
            kwargs = data['kwargs']
            if kwargs:
                output.append(f"                    assert_eq!(meta_line.arguments.kwargs.len(), {len(kwargs)});")
                for key, value in kwargs.items():
                    output.append(f"                    assert_eq!(meta_line.arguments.kwargs.iter().find(|(k, _)| k.content == {self._rust_string_literal(key)}).unwrap().1.content, {self._rust_string_literal(value)});")
            else:
                output.append("                    assert_eq!(meta_line.arguments.kwargs.len(), 0);")

        # VoidElement-specific: assert is_void is true
        output.append("                    assert!(meta_line.is_void, \"VoidElement should have is_void=true\");")

        output.append("                }")
        output.append(f"                _ => panic!(\"Expected VoidElement at index {index}, got {{:?}}\", events[{index}]),")
        output.append("            }")

        return output

    def _generate_text_finished_assertion(self, index: int, data: Dict[str, Any]) -> List[str]:
        """Generate assertions for TextFinished event.
        
        Only asserts on values explicitly provided in the YAML.
        """
        output = []
        
        # Handle shorthand: just an array of lines
        if isinstance(data, list):
            output.append(f"            match &events[{index}] {{")
            output.append("                Event::TextFinished { lines, links: _, footnotes: _ } => {")
            output.append(f"                    assert_eq!(lines.len(), {len(data)});")
            for i, line in enumerate(data):
                output.append(f"                    assert_eq!(lines[{i}], {self._rust_string_literal(line)});")
            output.append("                }")
            output.append(f"                _ => panic!(\"Expected TextFinished at index {index}, got {{:?}}\", events[{index}]),")
            output.append("            }")
            return output
        
        # Handle single line shorthand
        if isinstance(data, str):
            output.append(f"            match &events[{index}] {{")
            output.append("                Event::TextFinished { lines, links: _, footnotes: _ } => {")
            output.append(f"                    assert_eq!(lines.len(), 1);")
            output.append(f"                    assert_eq!(lines[0], {self._rust_string_literal(data)});")
            output.append("                }")
            output.append(f"                _ => panic!(\"Expected TextFinished at index {index}, got {{:?}}\", events[{index}]),")
            output.append("            }")
            return output
        
        output.append(f"            match &events[{index}] {{")
        output.append("                Event::TextFinished { lines, links, footnotes } => {")
        
        # Lines assertion (if present)
        if 'lines' in data:
            lines = data['lines']
            output.append(f"                    assert_eq!(lines.len(), {len(lines)});")
            for i, line in enumerate(lines):
                output.append(f"                    assert_eq!(lines[{i}], {self._rust_string_literal(line)});")
        
        # Links assertion (if present)
        if 'links' in data:
            links = data['links']
            output.append(f"                    // Checking {len(links)} links")
            # TODO: Implement full link assertions
        else:
            output.append("                    let _ = links;")
        
        # Footnotes assertion (if present)
        if 'footnotes' in data:
            footnotes = data['footnotes']
            output.append(f"                    assert_eq!(footnotes.len(), {len(footnotes)});")
            for i, footnote in enumerate(footnotes):
                if isinstance(footnote, dict):
                    if 'marker' in footnote:
                        output.append(f"                    assert_eq!(footnotes[{i}].marker, {footnote['marker']});")
                    if 'link' in footnote:
                        output.append(f"                    assert_eq!(footnotes[{i}].link, {self._rust_string_literal(footnote['link'])});")
        else:
            output.append("                    let _ = footnotes;")
        
        output.append("                }")
        output.append(f"                _ => panic!(\"Expected TextFinished at index {index}, got {{:?}}\", events[{index}]),")
        output.append("            }")
        
        return output
    
    def _rust_string_literal(self, s: str) -> str:
        """Convert string to Rust string literal."""
        # Escape backslashes and quotes
        escaped = s.replace('\\', '\\\\').replace('"', '\\"')
        # Handle newlines
        escaped = escaped.replace('\n', '\\n')
        return f'"{escaped}"'


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Generate Rust test files from Wasot syntax test markdown files."
    )
    parser.add_argument("input", help="Input markdown file with test cases")
    parser.add_argument("output", help="Output Rust test file")
    parser.add_argument(
        "--format", 
        action="store_true", 
        default=True,
        help="Auto-format and lint-fix generated code (default: enabled)"
    )
    parser.add_argument(
        "--no-format", 
        dest="format",
        action="store_false",
        help="Disable auto-formatting and linting"
    )
    
    args = parser.parse_args()
    
    input_path = Path(args.input)
    output_path = Path(args.output)
    
    # Read input file
    try:
        content = input_path.read_text()
    except Exception as e:
        print(f"Error reading {input_path}: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Parse test cases
    parser = MarkdownTestParser(content)
    test_cases = parser.parse()
    
    print(f"Parsed {len(test_cases)} test cases")
    
    # Generate Rust code
    generator = RustTestGenerator(test_cases)
    rust_code = generator.generate()
    
    # Write output file
    try:
        output_path.write_text(rust_code)
        print(f"Generated {output_path}")
    except Exception as e:
        print(f"Error writing {output_path}: {e}", file=sys.stderr)
        sys.exit(1)
    
    # Format and lint-fix the generated Rust file (if enabled)
    if args.format:
        try:
            print("Formatting generated code with rustfmt...")
            subprocess.run(
                ["rustfmt", str(output_path)],
                check=True,
                capture_output=True,
                text=True
            )
            print("✓ Formatted with rustfmt")
            
            print("Running clippy auto-fix on generated code...")
            subprocess.run(
                ["cargo", "clippy", "--fix", "--allow-dirty", "--allow-staged", "--", "--deny", "warnings"],
                cwd=output_path.parent.parent,  # Run from project root
                check=True,
                capture_output=True,
                text=True
            )
            print("✓ Applied clippy auto-fixes")
            
        except subprocess.CalledProcessError as e:
            print(f"Warning: Error during formatting/linting: {e}", file=sys.stderr)
            if e.stdout:
                print(f"stdout: {e.stdout}", file=sys.stderr)
            if e.stderr:
                print(f"stderr: {e.stderr}", file=sys.stderr)
            # Don't fail the script for formatting errors
        except FileNotFoundError as e:
            print(f"Warning: Tool not found: {e}", file=sys.stderr)
            print("Make sure rustfmt and cargo are installed and in PATH", file=sys.stderr)
    else:
        print("Skipping auto-formatting (disabled via --no-format)")


if __name__ == "__main__":
    main()