import os
import re
import requests
import argparse

def fix_markdown_images(md_file_path, output_file_path=None):
    base_url = "https://ocr.dolfprojects.com"
    base_assets_dir = "assets"
    
    # Set output file path to overwrite the original if not specified
    if not output_file_path:
        output_file_path = md_file_path

    try:
        with open(md_file_path, 'r', encoding='utf-8') as file:
            content = file.read()
    except FileNotFoundError:
        print(f"Error: The file '{md_file_path}' was not found.")
        return

    # Regex to find markdown images starting with /media/
    pattern = r'!\[([^\]]*)\]\((/media/[^\)]+)\)'
    matches = re.findall(pattern, content)

    if not matches:
        print("No images found to process.")
        return

    for alt_text, img_path in matches:
        full_url = base_url + img_path
        relative_path = img_path.lstrip('/')
        local_save_path = os.path.join(base_assets_dir, relative_path)
        new_md_path = f"./{base_assets_dir}{img_path}"

        os.makedirs(os.path.dirname(local_save_path), exist_ok=True)

        print(f"Downloading: {full_url}")
        try:
            response = requests.get(full_url, stream=True)
            response.raise_for_status() 
            
            with open(local_save_path, 'wb') as img_file:
                for chunk in response.iter_content(chunk_size=8192):
                    img_file.write(chunk)
            
            content = content.replace(f"({img_path})", f"({new_md_path})")
            print(f"Saved to: {local_save_path}")
            
        except requests.exceptions.RequestException as e:
            print(f"Failed to download {full_url}. Error: {e}")

    with open(output_file_path, 'w', encoding='utf-8') as file:
        file.write(content)
        
    print(f"\nSuccess! Updated Markdown saved to: {output_file_path}")

if __name__ == "__main__":
    # Set up argument parser
    parser = argparse.ArgumentParser(description="Fix image links in a Markdown file by downloading them locally.")
    
    # Add a required argument for the input file
    parser.add_argument("input_file", help="The path to the Markdown file you want to process.")
    
    # Add an optional argument for the output file
    parser.add_argument("-o", "--output", help="Optional: Path to save the new file. If omitted, the original file is overwritten.", default=None)
    
    args = parser.parse_args()
    
    # Run the function with the provided arguments
    fix_markdown_images(args.input_file, args.output)