import os
import sys
import json
import polib

def extract_untranslated(po_path, json_path):
    if not os.path.exists(po_path):
        print(f"Error: File not found: {po_path}")
        return False
        
    print(f"Loading PO file: {po_path}")
    po = polib.pofile(po_path)
    
    untranslated_data = {}
    
    for entry in po:
        if not entry.translated() or entry.fuzzy:
            # If it has a plural, extract both msgid and msgid_plural
            if entry.msgid_plural:
                if entry.msgid not in untranslated_data:
                    untranslated_data[entry.msgid] = ""
                if entry.msgid_plural not in untranslated_data:
                    untranslated_data[entry.msgid_plural] = ""
            else:
                if entry.msgid not in untranslated_data:
                    untranslated_data[entry.msgid] = ""
                    
    total_keys = len(untranslated_data)
    print(f"Extracted {total_keys} unique untranslated/fuzzy keys.")
    
    # Ensure output directory exists
    out_dir = os.path.dirname(json_path)
    if out_dir and not os.path.exists(out_dir):
        os.makedirs(out_dir)
        
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(untranslated_data, f, ensure_ascii=False)
        
    print(f"Saved JSON template to {json_path}")
    return True

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python extract_untranslated.py <po_path> <output_json_path>")
        sys.exit(1)
        
    po_file = sys.argv[1]
    json_file = sys.argv[2]
    extract_untranslated(po_file, json_file)