#!/usr/bin/env python3
"""
Script to add dimensions to remaining SHACL blocks that were missed in the first pass.
This handles:
1. Property blocks ending with ] . instead of ] ;
2. SPARQL constraint shapes
"""

import re
import sys
from pathlib import Path


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

    # 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 "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
    if "sh:node" in property_block:
        return "Conformity"

    return "Validity"


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

    # Pattern for property blocks ending with ] .
    property_pattern = r"(sh:property\s*\[)((?:[^\[\]]|\[(?:[^\[\]]|\[[^\[\]]*\])*\])*?)(\]\s*\.)"

    def replace_property(match):
        prefix = match.group(1)
        block_content = match.group(2)
        suffix = match.group(3)

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

        # Determine appropriate dimension
        dimension = determine_dimension(block_content)

        # Find sh:path line
        path_match = re.search(r"(sh:path\s+[^;]+;\s*\n)", block_content)
        if path_match:
            insert_pos = path_match.end()
            indent = "        "
            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

    content = re.sub(property_pattern, replace_property, content, flags=re.DOTALL)
    return content


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

    # Find shapes with sh:sparql that don't have dq:dimension
    # Pattern: NodeShape definition followed by properties, including sh:sparql
    shape_pattern = r"(^\w+:\w+\s*\n\s+a\s+sh:NodeShape\s*;[^\n]*\n(?:\s+[^\n]+\n)*?)(\s+sh:sparql\s*\[)"

    def replace_shape(match):
        shape_header = match.group(1)
        sparql_start = match.group(2)

        # Check if already has dq:dimension
        if "dq:dimension" in shape_header:
            return match.group(0)

        # Add dq:dimension before sh:sparql
        # Find the last property before sh:sparql
        lines = shape_header.split("\n")

        # Insert dimension before last line (which should be before sh:sparql)
        indent = "    "
        dimension_line = f'{indent}dq:dimension "Consistency" ;'

        # Insert after severity line or description line
        modified_lines = []
        inserted = False
        for i, line in enumerate(lines):
            modified_lines.append(line)
            if not inserted and ("sh:severity" in line or "sh:description" in line):
                # Insert after this line
                if i < len(lines) - 1:  # Not the last line
                    modified_lines.append(dimension_line)
                    inserted = True

        if not inserted and modified_lines:
            # Insert before the last line
            modified_lines.insert(-1, dimension_line)

        return "\n".join(modified_lines) + "\n" + sparql_start

    content = re.sub(shape_pattern, replace_shape, content, flags=re.MULTILINE)
    return content


def process_file(file_path: Path) -> tuple[bool, str]:
    """Process a single SHACL file to add remaining dimensions."""
    try:
        content = file_path.read_text(encoding="utf-8")
        original_content = content

        # Add dimensions to property blocks ending with period
        content = add_dimensions_to_property_blocks_ending_with_period(content)

        # Add dimensions to SPARQL shapes
        content = add_dimensions_to_sparql_shapes(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."""

    script_dir = Path(__file__).parent

    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 Remaining 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 remaining dimensions to SHACL rules!")
        return 0
    else:
        print("✓ All files already have dimensions.")
        return 0


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