#!/usr/bin/env python3
"""
Test tree-sitter-wasot parser against wasot parser.

Parses test files with both parsers and compares the structural output.
The comparison extracts key elements (block types, IDs, kwargs, etc.) from
both outputs and asserts they match.

Usage:
    python scripts/test_treesitter_parser.py                    # Run all tests
    python scripts/test_treesitter_parser.py -v                 # Verbose output
    python scripts/test_treesitter_parser.py --dir progressive_complexity
    python scripts/test_treesitter_parser.py --file tests/wasot_syntax_tests.md
"""

import argparse
import re
import subprocess
import sys
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional


@dataclass
class ParsedElement:
    """Represents a parsed element from either parser."""

    type: str
    level: int = 1
    id: Optional[str] = None
    is_void: bool = False
    is_context: bool = False
    kwargs: dict = field(default_factory=dict)
    children: list = field(default_factory=list)


def run_wasot_sexp(input_text: str) -> tuple[bool, str]:
    """Run wasot parser with sexp output."""
    try:
        result = subprocess.run(
            ["cargo", "run", "-q", "--bin", "wasot", "--", "-t", "sexp"],
            input=input_text,
            capture_output=True,
            text=True,
            timeout=10,
        )
        if result.returncode != 0:
            return False, result.stderr.strip()
        return True, result.stdout.strip()
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    except Exception as e:
        return False, str(e)


def run_treesitter_parse(input_text: str, ts_dir: Path) -> tuple[bool, str]:
    """Run tree-sitter parse on input text."""
    # Write to temp file since tree-sitter parse requires a file
    with tempfile.NamedTemporaryFile(mode="w", suffix=".wasot", delete=False) as f:
        f.write(input_text)
        temp_path = f.name

    try:
        result = subprocess.run(
            ["npx", "tree-sitter", "parse", temp_path],
            capture_output=True,
            text=True,
            timeout=10,
            cwd=ts_dir,
        )
        output = result.stdout.strip()
        # Check for ERROR or MISSING nodes
        has_error = "ERROR" in output or "MISSING" in output
        return not has_error, output
    except subprocess.TimeoutExpired:
        return False, "TIMEOUT"
    except Exception as e:
        return False, str(e)
    finally:
        Path(temp_path).unlink(missing_ok=True)


def pretty_print_sexp(sexp: str, indent: int = 2) -> str:
    """Pretty-print an s-expression with proper indentation."""
    result = []
    depth = 0
    i = 0
    in_string = False
    current_token = ""

    def flush_token():
        nonlocal current_token
        if current_token.strip():
            result.append(" " * (depth * indent) + current_token.strip())
        current_token = ""

    while i < len(sexp):
        char = sexp[i]

        # Handle string literals
        if char == '"' and (i == 0 or sexp[i-1] != '\\'):
            in_string = not in_string
            current_token += char
            i += 1
            continue

        if in_string:
            current_token += char
            i += 1
            continue

        if char == '(':
            flush_token()
            current_token = "("
            i += 1
            # Collect node name and any inline attributes
            while i < len(sexp) and sexp[i] != '(' and sexp[i] != ')':
                if sexp[i] == '"':
                    in_string = True
                    current_token += sexp[i]
                    i += 1
                    while i < len(sexp) and in_string:
                        current_token += sexp[i]
                        if sexp[i] == '"' and sexp[i-1] != '\\':
                            in_string = False
                        i += 1
                elif sexp[i] in '\n':
                    break
                else:
                    current_token += sexp[i]
                    i += 1
            depth += 1
        elif char == ')':
            if current_token.strip():
                result.append(" " * (depth * indent) + current_token.strip() + ")")
                current_token = ""
            else:
                if result:
                    result[-1] = result[-1] + ")"
            depth = max(0, depth - 1)
            i += 1
        elif char in ' \n\t':
            i += 1
        else:
            current_token += char
            i += 1

    flush_token()
    return "\n".join(result)


def extract_wasot_elements(sexp: str) -> list[dict]:
    """Extract structural elements from wasot sexp output."""
    elements = []

    # Pattern for quoted string that handles escaped quotes
    quoted_str = r'"(?:[^"\\]|\\.)*"'

    # Match block definitions: (block type: "..." level: N id: "..." ...)
    block_pattern = rf'\(block\s+type:\s+({quoted_str})\s+level:\s+(\d+)\s+id:\s+({quoted_str})'
    for match in re.finditer(block_pattern, sexp):
        elements.append(
            {
                "kind": "block",
                "type": match.group(1).strip('"'),
                "level": int(match.group(2)),
                "id": match.group(3).strip('"'),
            }
        )

    # Match void elements: (void type: "..." level: N id: "...")
    void_pattern = rf'\(void\s+type:\s+({quoted_str})\s+level:\s+(\d+)\s+id:\s+({quoted_str})'
    for match in re.finditer(void_pattern, sexp):
        elements.append(
            {
                "kind": "void",
                "type": match.group(1).strip('"'),
                "level": int(match.group(2)),
                "id": match.group(3).strip('"'),
            }
        )

    # Match context lines: (context-line type: "..." level: N ...)
    context_pattern = r'\(context-line\s+type:\s+"([^"]+)"\s+level:\s+(\d+)'
    for match in re.finditer(context_pattern, sexp):
        elements.append(
            {"kind": "context", "type": match.group(1), "level": int(match.group(2))}
        )

    # Match kwargs: (kwarg (key content: "..." ...) (value content: "..." ...))
    kwarg_pattern = r'\(kwarg\s+\(key\s+content:\s+"([^"]+)"[^)]*\)\s+\(value\s+content:\s+"([^"]+)"'
    for match in re.finditer(kwarg_pattern, sexp):
        elements.append({"kind": "kwarg", "key": match.group(1), "value": match.group(2)})

    # Match inline-text: (inline-text content: "..." ...)
    inline_pattern = r'\(inline-text\s+content:\s+"([^"]*)"'
    for match in re.finditer(inline_pattern, sexp):
        elements.append({"kind": "inline-text", "content": match.group(1)})

    return elements


def extract_treesitter_elements(output: str) -> list[dict]:
    """Extract structural elements from tree-sitter parse output."""
    elements = []

    # Find all meta_line blocks using regex on the full output
    # The output is indented, so we need to match across lines

    # First, normalize the output by joining lines
    # A meta_line starts with "(meta_line" and ends when parens balance

    # Simpler approach: find all meta_line occurrences and extract key info
    # Look for patterns in the full text

    # Find each meta_line section
    meta_line_pattern = r'\(meta_line \[\d+, \d+\] - \[\d+, \d+\]'

    # Split output into meta_line sections
    parts = re.split(r'(\(meta_line \[\d+, \d+\] - \[\d+, \d+\])', output)

    i = 1  # Start after first split (which is before first match)
    while i < len(parts):
        if '(meta_line' in parts[i]:
            # Collect content until parens balance
            meta_start = parts[i]
            content = parts[i + 1] if i + 1 < len(parts) else ""

            # Find where this meta_line ends by balancing parens
            full_meta = meta_start + content

            # Extract info from this meta_line
            element = parse_meta_line_content(full_meta)
            if element:
                elements.append(element)
        i += 2

    return elements


def parse_meta_line_content(content: str) -> Optional[dict]:
    """Parse content of a meta_line to extract element info."""
    # Count dots for level - look at the position span
    dots_match = re.search(r'\(dots \[(\d+), (\d+)\] - \[(\d+), (\d+)\]\)', content)
    if not dots_match:
        return None

    start_col = int(dots_match.group(2))
    end_col = int(dots_match.group(4))
    level = end_col - start_col  # Number of dots = level

    # Check for context line with content (./type)
    if '(context_content' in content:
        return {"kind": "context", "level": level}

    # Check for close marker (./) - not a block or context
    if '(context_marker' in content:
        return None

    # Check for void (has void_marker)
    is_void = '(void_marker' in content

    # Check for block_id
    has_id = '(block_id' in content

    # Check for kwarg
    has_kwarg = '(kwarg' in content

    # A meta_line with just dots (no block_type) is still a block with empty type
    return {
        "kind": "void" if is_void else "block",
        "level": level,
        "has_id": has_id,
        "has_kwarg": has_kwarg,
    }


def compare_elements(
    wasot_elements: list[dict], ts_elements: list[dict]
) -> tuple[bool, list[str]]:
    """Compare elements from both parsers."""
    differences = []

    # Compare counts by kind
    wasot_blocks = [e for e in wasot_elements if e["kind"] == "block"]
    ts_blocks = [e for e in ts_elements if e["kind"] == "block"]

    wasot_voids = [e for e in wasot_elements if e["kind"] == "void"]
    ts_voids = [e for e in ts_elements if e["kind"] == "void"]

    wasot_contexts = [e for e in wasot_elements if e["kind"] == "context"]
    ts_contexts = [e for e in ts_elements if e["kind"] == "context"]

    if len(wasot_blocks) != len(ts_blocks):
        differences.append(
            f"Block count mismatch: wasot={len(wasot_blocks)}, ts={len(ts_blocks)}"
        )

    if len(wasot_voids) != len(ts_voids):
        differences.append(
            f"Void count mismatch: wasot={len(wasot_voids)}, ts={len(ts_voids)}"
        )

    if len(wasot_contexts) != len(ts_contexts):
        differences.append(
            f"Context count mismatch: wasot={len(wasot_contexts)}, ts={len(ts_contexts)}"
        )

    # Compare levels for blocks
    wasot_levels = sorted([e["level"] for e in wasot_blocks])
    ts_levels = sorted([e["level"] for e in ts_blocks])

    if wasot_levels != ts_levels:
        differences.append(f"Block levels mismatch: wasot={wasot_levels}, ts={ts_levels}")

    return len(differences) == 0, differences


@dataclass
class TestResult:
    """Result of a single test."""

    name: str
    passed: bool
    wasot_ok: bool
    ts_ok: bool
    differences: list[str] = field(default_factory=list)
    wasot_output: str = ""
    ts_output: str = ""


def run_test(
    name: str, input_text: str, ts_dir: Path, verbose: bool = False
) -> TestResult:
    """Run a single test comparing both parsers."""
    result = TestResult(name=name, passed=False, wasot_ok=False, ts_ok=False)

    # Run wasot
    wasot_ok, wasot_output = run_wasot_sexp(input_text)
    result.wasot_ok = wasot_ok
    result.wasot_output = wasot_output

    # Run tree-sitter
    ts_ok, ts_output = run_treesitter_parse(input_text, ts_dir)
    result.ts_ok = ts_ok
    result.ts_output = ts_output

    if not wasot_ok:
        result.differences.append(f"Wasot parse failed: {wasot_output}")
        return result

    if not ts_ok:
        result.differences.append(f"Tree-sitter parse failed or has errors")
        return result

    # Extract and compare elements
    wasot_elements = extract_wasot_elements(wasot_output)
    ts_elements = extract_treesitter_elements(ts_output)

    passed, differences = compare_elements(wasot_elements, ts_elements)
    result.passed = passed
    result.differences = differences

    return result


def load_progressive_complexity_tests(
    prog_dir: Path,
) -> list[tuple[str, str]]:
    """Load test cases from progressive_complexity directory."""
    tests = []
    for file_path in sorted(prog_dir.glob("*.wasot")):
        content = file_path.read_text()
        tests.append((file_path.name, content))
    return tests


def load_markdown_tests(md_path: Path) -> list[tuple[str, str]]:
    """Load test cases from markdown file."""
    # Import the parser from parse_syntax_file
    sys.path.insert(0, str(md_path.parent.parent / "scripts"))
    try:
        from parse_syntax_file import MarkdownTestParser

        content = md_path.read_text()
        parser = MarkdownTestParser(content)
        test_cases = parser.parse()
        return [(tc.name, tc.input) for tc in test_cases]
    except ImportError:
        print("Warning: Could not import MarkdownTestParser", file=sys.stderr)
        return []


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Test tree-sitter-wasot parser against wasot parser."
    )
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="Show verbose output"
    )
    parser.add_argument(
        "--dir",
        type=Path,
        default=None,
        help="Directory with .wasot test files (default: progressive_complexity)",
    )
    parser.add_argument(
        "--file", type=Path, default=None, help="Markdown file with test cases"
    )
    parser.add_argument(
        "--stop-on-fail", action="store_true", help="Stop on first failure"
    )
    parser.add_argument(
        "--show-sexp", action="store_true", help="Show pretty-printed s-expressions from both parsers"
    )
    parser.add_argument(
        "--test", type=str, default=None, help="Run only this specific test file"
    )

    args = parser.parse_args()

    # Find paths
    script_dir = Path(__file__).parent
    project_root = script_dir.parent
    ts_dir = project_root / "tree-sitter-wasot"
    prog_dir = project_root / "progressive_complexity"

    if not ts_dir.exists():
        print(f"Error: {ts_dir} does not exist", file=sys.stderr)
        sys.exit(1)

    # Load test cases
    tests: list[tuple[str, str]] = []

    if args.file:
        tests = load_markdown_tests(args.file)
    elif args.dir:
        tests = load_progressive_complexity_tests(args.dir)
    else:
        # Default: use progressive_complexity
        if prog_dir.exists():
            tests = load_progressive_complexity_tests(prog_dir)
        else:
            print(f"Error: No test directory found", file=sys.stderr)
            sys.exit(1)

    if not tests:
        print("No tests found", file=sys.stderr)
        sys.exit(1)

    print(f"Running {len(tests)} tests...\n")

    # Run tests
    results: list[TestResult] = []
    passed_count = 0
    failed_count = 0

    for name, input_text in tests:
        # Filter by specific test if requested
        if args.test and args.test not in name:
            continue

        result = run_test(name, input_text, ts_dir, args.verbose)
        results.append(result)

        if result.passed:
            passed_count += 1
            if args.verbose:
                print(f"✓ {name}")
        else:
            failed_count += 1
            print(f"✗ {name}")
            if result.differences:
                for diff in result.differences:
                    print(f"  - {diff}")
            if args.verbose:
                print(f"  Wasot output: {result.wasot_output[:200]}...")
                print(f"  TS output: {result.ts_output[:200]}...")

            if args.stop_on_fail:
                break

        # Show pretty s-expressions if requested
        if args.show_sexp:
            print(f"\n{'=' * 60}")
            print(f"TEST: {name}")
            print(f"{'=' * 60}")
            print(f"\nInput:")
            print("-" * 40)
            for line in input_text.split('\n'):
                print(f"  {repr(line)}")
            print("-" * 40)
            print(f"\nWasot S-expression:")
            print("-" * 40)
            print(pretty_print_sexp(result.wasot_output))
            print("-" * 40)
            print(f"\nTree-sitter S-expression:")
            print("-" * 40)
            print(result.ts_output)
            print("-" * 40)

    # Summary
    print(f"\n{'=' * 60}")
    print(f"SUMMARY")
    print(f"{'=' * 60}")
    print(f"Total:  {len(results)}")
    print(f"Passed: {passed_count}")
    print(f"Failed: {failed_count}")

    if failed_count > 0:
        print(f"\nFailed tests:")
        for r in results:
            if not r.passed:
                print(f"  - {r.name}")
        sys.exit(1)
    else:
        print(f"\n✓ All tests passed!")
        sys.exit(0)


if __name__ == "__main__":
    main()
