#!/usr/bin/env python3
"""Restore correct code block languages by comparing with archived originals."""

import re
import sys
from pathlib import Path

def extract_code_blocks(content):
    """Extract all code blocks with line numbers and languages."""
    blocks = []
    lines = content.split('\n')
    in_block = False
    block_start = 0
    block_lang = ''

    for i, line in enumerate(lines):
        if line.startswith('```'):
            if not in_block:
                # Opening fence
                block_start = i
                block_lang = line[3:].strip()
                in_block = True
            else:
                # Closing fence
                blocks.append({
                    'start': block_start,
                    'end': i,
                    'language': block_lang,
                    'content': '\n'.join(lines[block_start+1:i])
                })
                in_block = False

    return blocks

def find_matching_block(current_content, original_blocks, block_index):
    """Find matching block in current content by comparing content."""
    current_blocks = extract_code_blocks(current_content)

    if block_index < len(current_blocks) and block_index < len(original_blocks):
        current = current_blocks[block_index]
        original = original_blocks[block_index]

        # Compare content to verify it's the same block
        if current['content'].strip() == original['content'].strip():
            return original['language']

    return None

def restore_languages(current_content, original_content):
    """Restore original code block languages."""
    original_blocks = extract_code_blocks(original_content)

    if not original_blocks:
        return current_content

    lines = current_content.split('\n')
    in_block = False
    block_index = 0
    block_start = 0

    for i, line in enumerate(lines):
        if line.startswith('```'):
            if not in_block:
                # Opening fence - restore original language if we have it
                if block_index < len(original_blocks):
                    original_lang = original_blocks[block_index]['language']
                    if original_lang:
                        lines[i] = '```' + original_lang
                    block_start = i
                in_block = True
            else:
                # Closing fence - keep as just ```
                lines[i] = '```'
                block_index += 1
                in_block = False

    return '\n'.join(lines)

def process_file(current_path, original_path):
    """Process a single file pair."""
    try:
        with open(current_path, 'r', encoding='utf-8') as f:
            current_content = f.read()

        if not original_path.exists():
            print(f"⚠ Original not found: {original_path.name}")
            return False

        with open(original_path, 'r', encoding='utf-8') as f:
            original_content = f.read()

        fixed_content = restore_languages(current_content, original_content)

        with open(current_path, 'w', encoding='utf-8') as f:
            f.write(fixed_content)

        return True
    except Exception as e:
        print(f"✗ Error: {e}", file=sys.stderr)
        return False

def main():
    """Restore correct code block languages."""
    current_root = Path('provisioning/docs/src')
    original_root = Path('.coder/archive/docs-pre-audit/stubs')

    # All AI files to restore
    ai_files = [
        'ai/ai-agents.md',
        'ai/ai-assisted-forms.md',
        'ai/architecture.md',
        'ai/config-generation.md',
        'ai/configuration.md',
        'ai/cost-management.md',
        'ai/mcp-integration.md',
        'ai/natural-language-config.md',
        'ai/rag-system.md',
        'ai/README.md',
        'ai/security-policies.md',
        'ai/troubleshooting-with-ai.md',
    ]

    success_count = 0
    for filepath_rel in ai_files:
        current_path = current_root / filepath_rel
        original_path = original_root / filepath_rel

        if current_path.exists() and original_path.exists():
            if process_file(current_path, original_path):
                print(f"✓ Restored {filepath_rel}")
                success_count += 1
            else:
                print(f"✗ Failed to restore {filepath_rel}")
        else:
            if not current_path.exists():
                print(f"⚠ Current not found: {filepath_rel}")
            if not original_path.exists():
                print(f"⚠ Original not found: {filepath_rel}")

    print(f"\n✓ Restored {success_count}/{len(ai_files)} files")
    return 0 if success_count == len(ai_files) else 1

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