#!/usr/bin/env python3

import sys
import os
import argparse
import re
import json
import base64
import zlib
from pathlib import Path

def pako_deflate(data):
    """Compress data using the same parameters as pako.deflate"""
    compress = zlib.compressobj(9, zlib.DEFLATED, 15, 8, zlib.Z_DEFAULT_STRATEGY)
    compressed_data = compress.compress(data)
    compressed_data += compress.flush()
    return compressed_data

def generate_mermaid_url(mermaid_code):
    """Generate a mermaid.live URL from mermaid code"""
    # Create the mermaid object structure expected by mermaid.live
    mermaid_obj = {
        "code": mermaid_code.strip(),
        "mermaid": {
            "theme": "default"
        }
    }
    
    # Convert to JSON
    mermaid_json = json.dumps(mermaid_obj)
    
    # Compress using pako deflate parameters
    byte_data = mermaid_json.encode('ascii')
    compressed = pako_deflate(byte_data)
    
    # Base64 encode (URL-safe)
    encoded = base64.b64encode(compressed).decode('ascii')
    encoded = encoded.replace('+', '-').replace('/', '_')
    
    return f"https://mermaid.live/edit#pako:{encoded}"

def process_markdown_file(file_path, dry_run=False):
    """Process a single markdown file to add mermaid.live links"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"Error reading {file_path}: {e}", file=sys.stderr)
        return False
    
    # Pattern to match mermaid code blocks
    # First find all code blocks, then check if they contain mermaid
    code_block_pattern = r'```(\w*)\n(.*?)\n```'
    
    # Find all code blocks and check if they're mermaid
    matches = list(re.finditer(code_block_pattern, content, re.MULTILINE | re.DOTALL))
    
    # Filter for mermaid blocks
    mermaid_matches = []
    mermaid_keywords = ['graph', 'flowchart', 'sequenceDiagram', 'classDiagram', 'stateDiagram', 
                       'erDiagram', 'journey', 'gantt', 'pie', 'gitgraph', 'mindmap', 
                       'timeline', 'quadrantChart', 'xyChart', 'block', 'architecture']
    
    for match in matches:
        language = match.group(1).lower()  # Language specifier after ```
        code_content = match.group(2)     # Code content
        
        # Check if it's a mermaid block
        is_mermaid = (language == 'mermaid' or 
                     any(keyword in code_content for keyword in mermaid_keywords))
        
        if is_mermaid:
            mermaid_matches.append(match)
    
    if not mermaid_matches:
        return False
    
    # Process matches in reverse order to maintain correct positions
    modified_content = content
    changes_made = False
    
    for match in reversed(mermaid_matches):
        language = match.group(1)         # Language specifier 
        mermaid_code = match.group(2)     # The actual mermaid code
        
        # Check if there's already a link before this mermaid block
        text_before = modified_content[:match.start()]
        link_pattern = r'\[link to rendered version here\]\(https://mermaid\.live/[^)]+\)\s*$'
        
        if re.search(link_pattern, text_before):
            continue  # Skip if link already exists
        
        try:
            # Generate the mermaid.live URL
            mermaid_url = generate_mermaid_url(mermaid_code)
            
            # Create the link text
            link_text = f"[link to rendered version here]({mermaid_url})\n\n"
            
            # Insert the link before the mermaid block
            before = modified_content[:match.start()]
            after = modified_content[match.start():]
            modified_content = before + link_text + after
            changes_made = True
            
            print(f"Added link for mermaid diagram in {file_path}")
            
        except Exception as e:
            print(f"Error generating URL for mermaid diagram in {file_path}: {e}", file=sys.stderr)
            continue
    
    # Write the modified content back to the file
    if changes_made and not dry_run:
        try:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(modified_content)
            print(f"Updated {file_path}")
        except Exception as e:
            print(f"Error writing {file_path}: {e}", file=sys.stderr)
            return False
    elif changes_made and dry_run:
        print(f"[DRY RUN] Would update {file_path}")
    
    return changes_made

def main():
    parser = argparse.ArgumentParser(
        description="Add mermaid.live links to mermaid diagrams in markdown files",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 add-mermaid-links.py --target-dir ./docs
  python3 add-mermaid-links.py --target-dir ./docs --dry-run
  python3 add-mermaid-links.py --target-dir ./docs --recursive
        """
    )
    
    parser.add_argument(
        '--target-dir',
        required=True,
        help='Directory to scan for markdown files'
    )
    
    parser.add_argument(
        '--recursive', '-r',
        action='store_true',
        help='Recursively scan subdirectories'
    )
    
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Show what would be changed without modifying files'
    )
    
    args = parser.parse_args()
    
    target_dir = Path(args.target_dir)
    
    if not target_dir.exists():
        print(f"Error: Directory {target_dir} does not exist", file=sys.stderr)
        sys.exit(1)
    
    if not target_dir.is_dir():
        print(f"Error: {target_dir} is not a directory", file=sys.stderr)
        sys.exit(1)
    
    # Find all markdown files
    if args.recursive:
        md_files = list(target_dir.rglob("*.md"))
    else:
        md_files = list(target_dir.glob("*.md"))
    
    if not md_files:
        print(f"No markdown files found in {target_dir}")
        sys.exit(0)
    
    print(f"Found {len(md_files)} markdown files")
    if args.dry_run:
        print("DRY RUN MODE - No files will be modified")
    
    # Process each file
    files_modified = 0
    for md_file in md_files:
        if process_markdown_file(md_file, args.dry_run):
            files_modified += 1
    
    print(f"Processed {len(md_files)} files, {files_modified} files {'would be ' if args.dry_run else ''}modified")

if __name__ == "__main__":
    main()