#!/usr/bin/env python3
"""
Script to add data quality dimensions to SHACL rule files.

This script analyzes SHACL property constraints and adds appropriate
dq:dimension annotations based on the constraint types:
- Completeness: Required fields (minCount/maxCount)
- Validity: Format/type validation (pattern, datatype, minLength, maxLength)
- Conformity: Controlled vocabularies (sh:in) and standards compliance
- Consistency: Business rules and cross-property validation (sh:sparql, sh:node for relations)
"""

import re
import sys
from pathlib import Path


def has_dq_prefix(content: str) -> bool:
    """Check if the file already has the dq: prefix defined."""
    return bool(re.search(r"@prefix\s+dq:\s*<[^>]+>", content))


def add_dq_prefix(content: str) -> str:
    """Add the dq: prefix to the file if not present."""
    if has_dq_prefix(content):
        return content

    # Find the last @prefix line and add dq: after it
    prefix_lines = list(re.finditer(r"^@prefix\s+[^.]+\.\s*$", content, re.MULTILINE))
    if prefix_lines:
        last_prefix = prefix_lines[-1]
        insert_pos = last_prefix.end()
        dq_prefix = "\n@prefix dq: <http://example.org/dataquality/> ."
        content = content[:insert_pos] + dq_prefix + content[insert_pos:]

    return content


def determine_dimension(property_block: str) -> str:
    """
    Determine the appropriate data quality dimension based on constraints.

    Priority order:
    1. sh:in -> Conformity (controlled vocabulary)
    2. sh:minCount (with high priority) -> Completeness (required field)
    3. sh:pattern, sh:datatype with format validation -> Validity
    4. sh:node -> Conformity (referential integrity/shape conformance)
    5. Default for other cases -> Validity
    """

    # Controlled vocabulary (highest priority for property constraints)
    if "sh:in" in property_block:
        return "Conformity"

    # Required field with minCount
    if "sh:minCount" in property_block and "sh:minCount 1" in property_block:
        # If it also has sh:node, it's checking both completeness and conformity
        # Primary concern is completeness for required fields
        if "sh:node" in property_block:
            return "Completeness"
        return "Completeness"

    # Format/pattern validation
    if "sh:pattern" in property_block:
        return "Validity"

    # Type validation with constraints
    if "sh:datatype" in property_block and any(
        x in property_block for x in ["sh:minLength", "sh:maxLength", "sh:minInclusive", "sh:maxInclusive"]
    ):
        return "Validity"

    # Shape conformance / referential integrity (no minCount, so optional)
    if "sh:node" in property_block:
        return "Conformity"

    # Default for other property constraints
    return "Validity"


def add_dimensions_to_property_blocks(content: str) -> str:
    """Add dq:dimension to sh:property blocks that don't have it."""

    # Pattern to match sh:property blocks
    # This regex captures the entire property block including nested brackets
    property_pattern = r"(sh:property\s*\[)((?:[^\[\]]|\[(?:[^\[\]]|\[[^\[\]]*\])*\])*?)(\]\s*;)"

    def replace_property(match):
        prefix = match.group(1)  # 'sh:property ['
        block_content = match.group(2)  # everything inside
        suffix = match.group(3)  # '] ;'

        # Check if this block already has dq:dimension
        if "dq:dimension" in block_content:
            return match.group(0)  # Return unchanged

        # Determine appropriate dimension
        dimension = determine_dimension(block_content)

        # Find the first sh: constraint after sh:path to insert dimension
        # Look for sh:path line
        path_match = re.search(r"(sh:path\s+[^;]+;\s*\n)", block_content)
        if path_match:
            insert_pos = path_match.end()
            # Insert dimension with proper indentation
            indent = "        "  # 8 spaces to match typical indentation
            dimension_line = f'{indent}dq:dimension "{dimension}" ;\n'
            block_content = block_content[:insert_pos] + dimension_line + block_content[insert_pos:]

        return prefix + block_content + suffix

    # Replace all property blocks
    content = re.sub(property_pattern, replace_property, content, flags=re.DOTALL)

    return content


def add_dimensions_to_sparql_blocks(content: str) -> str:
    """Add dq:dimension to shapes with sh:sparql constraints."""

    # Pattern to match NodeShapes with sparql constraints
    # We'll add dimension at the shape level for SPARQL rules

    def replace_sparql_shape(match):
        shape_text = match.group(0)

        # Check if already has dq:dimension
        if "dq:dimension" in shape_text:
            return shape_text

        # For SPARQL constraints, typically these are consistency checks
        # Insert dq:dimension before sh:sparql
        insertion_point = shape_text.rfind("sh:sparql")
        if insertion_point > 0:
            # Find the line break before sh:sparql
            line_start = shape_text.rfind("\n", 0, insertion_point)
            indent = "    "  # 4 spaces for shape-level properties
            dimension_line = f'\n{indent}dq:dimension "Consistency" ;'
            shape_text = shape_text[:line_start] + dimension_line + shape_text[line_start:]

        return shape_text

    # Note: SPARQL shapes are tricky. For now, skip automatic addition to SPARQL shapes
    # as they may need manual review
    return content


def process_file(file_path: Path) -> tuple[bool, str]:
    """
    Process a single SHACL file to add dimensions.

    Returns:
        (success, message) tuple
    """
    try:
        content = file_path.read_text(encoding="utf-8")
        original_content = content

        # Add dq: prefix if not present
        content = add_dq_prefix(content)

        # Add dimensions to property blocks
        content = add_dimensions_to_property_blocks(content)

        # Skip SPARQL blocks for now - they need manual review
        # content = add_dimensions_to_sparql_blocks(content)

        # Only write if content changed
        if content != original_content:
            file_path.write_text(content, encoding="utf-8")
            return True, f"✓ Updated: {file_path.name}"
        else:
            return True, f"- Skipped (no changes needed): {file_path.name}"

    except Exception as e:
        return False, f"✗ Error processing {file_path.name}: {e!s}"


def main():
    """Main function to process all SHACL files."""

    # Get the script directory
    script_dir = Path(__file__).parent

    # Define paths to process (excluding .venv and test_templates)
    environments = [
        script_dir / "config/environments/abp-dev/shacl_rules",
        script_dir / "config/environments/cog-ai/shacl_rules",
        script_dir / "config/environments/cog-sail/shacl_rules",
    ]

    print("=" * 80)
    print("Adding Data Quality Dimensions to SHACL Rules")
    print("=" * 80)
    print()

    total_files = 0
    updated_files = 0
    skipped_files = 0
    error_files = 0

    for env_path in environments:
        if not env_path.exists():
            print(f"Warning: Path not found: {env_path}")
            continue

        print(f"\nProcessing: {env_path.relative_to(script_dir)}")
        print("-" * 80)

        ttl_files = sorted(env_path.glob("*.ttl"))

        for ttl_file in ttl_files:
            total_files += 1
            success, message = process_file(ttl_file)
            print(f"  {message}")

            if success:
                if "Updated" in message:
                    updated_files += 1
                else:
                    skipped_files += 1
            else:
                error_files += 1

    # Summary
    print()
    print("=" * 80)
    print("Summary")
    print("=" * 80)
    print(f"Total files processed: {total_files}")
    print(f"Files updated: {updated_files}")
    print(f"Files skipped: {skipped_files}")
    print(f"Errors: {error_files}")
    print()

    if error_files > 0:
        print("⚠️  Some files had errors. Please review the output above.")
        return 1
    elif updated_files > 0:
        print("✓ Successfully added dimensions to SHACL rules!")
        print()
        print("Note: SPARQL constraints were not automatically annotated.")
        print("Please review files with sh:sparql blocks and add dimensions manually.")
        return 0
    else:
        print("✓ All files already have dimensions or no changes needed.")
        return 0


if __name__ == "__main__":
    sys.exit(main())
