#!/usr/bin/env python3
"""
Extract Cypher queries from TCK feature files and libcypher-parser tests.

This creates a comprehensive test suite for validating grammar completeness.
"""

import re
import json
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Optional


@dataclass
class TestQuery:
    query: str
    source_file: str
    source_line: int
    scenario: Optional[str] = None
    category: Optional[str] = None
    expected_error: bool = False


def extract_from_feature_file(filepath: Path) -> list[TestQuery]:
    """Extract queries from a Gherkin .feature file."""
    queries = []
    content = filepath.read_text()

    # Determine category from path
    parts = filepath.parts
    if 'clauses' in parts:
        idx = parts.index('clauses')
        category = '/'.join(parts[idx+1:-1]) if idx + 1 < len(parts) - 1 else parts[-2]
    elif 'expressions' in parts:
        idx = parts.index('expressions')
        category = '/'.join(parts[idx+1:-1]) if idx + 1 < len(parts) - 1 else parts[-2]
    else:
        category = filepath.parent.name

    current_scenario = None
    in_query_block = False
    query_lines = []

    lines = content.split('\n')
    for i, line in enumerate(lines, 1):
        stripped = line.strip()

        # Track scenario names
        if stripped.startswith('Scenario:'):
            current_scenario = stripped[9:].strip()
        elif stripped.startswith('Scenario Outline:'):
            current_scenario = stripped[17:].strip()

        # Detect query blocks
        if stripped.startswith('When executing query:') or \
           stripped.startswith('And executing query:') or \
           stripped.startswith('When executing control query:'):
            in_query_block = True
            query_lines = []
            continue

        # Detect query in Given clause
        if 'having executed:' in stripped.lower():
            in_query_block = True
            query_lines = []
            continue

        # End of query block
        if in_query_block:
            if stripped.startswith('"""'):
                if query_lines:
                    # End of query
                    query_text = '\n'.join(query_lines).strip()
                    if query_text:
                        queries.append(TestQuery(
                            query=query_text,
                            source_file=str(filepath),
                            source_line=i,
                            scenario=current_scenario,
                            category=category,
                        ))
                    in_query_block = False
                    query_lines = []
                continue
            elif stripped and not stripped.startswith('#'):
                query_lines.append(stripped)

        # Single-line query patterns
        query_match = re.search(r'executing query:\s*`([^`]+)`', stripped)
        if query_match:
            queries.append(TestQuery(
                query=query_match.group(1),
                source_file=str(filepath),
                source_line=i,
                scenario=current_scenario,
                category=category,
            ))

    return queries


def extract_from_libcypher_test(filepath: Path) -> list[TestQuery]:
    """Extract queries from libcypher-parser C test files."""
    queries = []
    content = filepath.read_text()

    # Pattern for cypher_parse("query") or cypher_parse_result(..., "query", ...)
    patterns = [
        r'cypher_parse\s*\(\s*"([^"]+)"',
        r'cypher_parse_result\s*\([^,]+,\s*"([^"]+)"',
        r'result\s*=\s*parse\s*\(\s*"([^"]+)"',
    ]

    for pattern in patterns:
        for i, match in enumerate(re.finditer(pattern, content)):
            query = match.group(1)
            # Unescape C string escapes
            query = query.replace('\\n', '\n')
            query = query.replace('\\t', '\t')
            query = query.replace('\\"', '"')
            query = query.replace('\\\\', '\\')

            queries.append(TestQuery(
                query=query,
                source_file=str(filepath),
                source_line=content[:match.start()].count('\n') + 1,
                category='libcypher-parser',
            ))

    return queries


def categorize_query(query: str) -> list[str]:
    """Identify what Cypher features a query uses."""
    features = []
    query_upper = query.upper()

    # Clauses
    clause_patterns = [
        ('MATCH', 'match'),
        ('OPTIONAL MATCH', 'optional-match'),
        ('WHERE', 'where'),
        ('RETURN', 'return'),
        ('WITH', 'with'),
        ('UNWIND', 'unwind'),
        ('CREATE', 'create'),
        ('MERGE', 'merge'),
        ('DELETE', 'delete'),
        ('DETACH DELETE', 'detach-delete'),
        ('SET', 'set'),
        ('REMOVE', 'remove'),
        ('FOREACH', 'foreach'),
        ('CALL', 'call'),
        ('YIELD', 'yield'),
        ('UNION', 'union'),
        ('ORDER BY', 'order-by'),
        ('SKIP', 'skip'),
        ('LIMIT', 'limit'),
        ('DISTINCT', 'distinct'),
    ]

    for pattern, feature in clause_patterns:
        if pattern in query_upper:
            features.append(feature)

    # Expression features
    expr_patterns = [
        (r'\bCASE\b', 'case-expression'),
        (r'\bWHEN\b.*\bTHEN\b', 'case-when'),
        (r'\bIN\b\s*\[', 'list-predicate'),
        (r'\bSTARTS WITH\b', 'starts-with'),
        (r'\bENDS WITH\b', 'ends-with'),
        (r'\bCONTAINS\b', 'contains'),
        (r'\bIS NULL\b', 'null-check'),
        (r'\bIS NOT NULL\b', 'null-check'),
        (r'\[\s*\w+\s+IN\b', 'list-comprehension'),
        (r'\bALL\s*\(', 'all-predicate'),
        (r'\bANY\s*\(', 'any-predicate'),
        (r'\bNONE\s*\(', 'none-predicate'),
        (r'\bSINGLE\s*\(', 'single-predicate'),
        (r'\bEXISTS\s*\{', 'exists-subquery'),
        (r'\bCOUNT\s*\{', 'count-subquery'),
        (r'-\[.*\*', 'variable-length-path'),
        (r'shortestPath\s*\(', 'shortest-path'),
        (r'allShortestPaths\s*\(', 'all-shortest-paths'),
    ]

    for pattern, feature in expr_patterns:
        if re.search(pattern, query, re.IGNORECASE):
            features.append(feature)

    # Aggregation functions
    agg_funcs = ['count', 'sum', 'avg', 'min', 'max', 'collect', 'stdev', 'percentile']
    for func in agg_funcs:
        if re.search(rf'\b{func}\s*\(', query, re.IGNORECASE):
            features.append(f'agg-{func}')

    return features


def main():
    project_root = Path(__file__).parent.parent
    output_dir = Path(__file__).parent / 'tck-queries'
    output_dir.mkdir(exist_ok=True)

    all_queries = []

    # Extract from TCK features
    tck_dir = project_root / 'tck_features'
    if tck_dir.exists():
        print(f"Extracting from TCK features in {tck_dir}...")
        feature_files = list(tck_dir.rglob('*.feature'))
        print(f"  Found {len(feature_files)} feature files")

        for fpath in feature_files:
            queries = extract_from_feature_file(fpath)
            all_queries.extend(queries)

        print(f"  Extracted {len(all_queries)} queries from TCK")

    # Extract from libcypher-parser tests
    libcypher_dir = Path('/Users/gregoriomomm/libcypher-parser/lib/test')
    if libcypher_dir.exists():
        print(f"\nExtracting from libcypher-parser tests in {libcypher_dir}...")
        test_files = list(libcypher_dir.glob('check_*.c'))
        print(f"  Found {len(test_files)} test files")

        initial_count = len(all_queries)
        for fpath in test_files:
            queries = extract_from_libcypher_test(fpath)
            all_queries.extend(queries)

        print(f"  Extracted {len(all_queries) - initial_count} queries from libcypher-parser")

    # Deduplicate by query text
    seen = set()
    unique_queries = []
    for q in all_queries:
        normalized = ' '.join(q.query.split())
        if normalized not in seen:
            seen.add(normalized)
            unique_queries.append(q)

    print(f"\nTotal unique queries: {len(unique_queries)}")

    # Categorize queries
    feature_counts = {}
    for q in unique_queries:
        features = categorize_query(q.query)
        for f in features:
            feature_counts[f] = feature_counts.get(f, 0) + 1

    # Save queries to JSON
    output_file = output_dir / 'all_queries.json'
    with open(output_file, 'w') as f:
        json.dump([asdict(q) for q in unique_queries], f, indent=2)
    print(f"\nSaved queries to {output_file}")

    # Save just the query texts for simple parsing tests
    queries_only = output_dir / 'queries.txt'
    with open(queries_only, 'w') as f:
        for q in unique_queries:
            # Write as single line with escaped newlines
            query_line = q.query.replace('\n', ' ').strip()
            f.write(query_line + '\n')
    print(f"Saved query texts to {queries_only}")

    # Save feature summary
    summary_file = output_dir / 'feature_summary.json'
    with open(summary_file, 'w') as f:
        json.dump({
            'total_queries': len(unique_queries),
            'feature_counts': dict(sorted(feature_counts.items(), key=lambda x: -x[1])),
            'sources': {
                'tck_features': sum(1 for q in unique_queries if 'tck_features' in q.source_file),
                'libcypher': sum(1 for q in unique_queries if 'libcypher' in q.source_file),
            }
        }, f, indent=2)
    print(f"Saved feature summary to {summary_file}")

    # Print top features
    print("\nTop features by query count:")
    for feature, count in sorted(feature_counts.items(), key=lambda x: -x[1])[:20]:
        print(f"  {feature}: {count}")


if __name__ == '__main__':
    main()
