import os
import sys
import json
import polib

def apply_translations(po_path, json_path):
    if not os.path.exists(po_path):
        print(f"Error: PO file not found: {po_path}")
        return False
    if not os.path.exists(json_path):
        print(f"Error: JSON translation file not found: {json_path}")
        return False
        
    print(f"Loading PO file: {po_path}")
    po = polib.pofile(po_path)
    
    print(f"Loading JSON translations: {json_path}")
    with open(json_path, 'r', encoding='utf-8') as f:
        translations = json.load(f)
        
    applied_count = 0
    
    for entry in po:
        if not entry.translated() or entry.fuzzy:
            updated = False
            
            # Handle plural entries
            if entry.msgid_plural:
                # Singular part (index 0)
                if entry.msgid in translations:
                    val = translations[entry.msgid]
                    if val and val.strip():
                        entry.msgstr_plural[0] = val
                        updated = True
                
                # Plural part (index 1, and higher indices if present)
                if entry.msgid_plural in translations:
                    val = translations[entry.msgid_plural]
                    if val and val.strip():
                        for plural_index in entry.msgstr_plural.keys():
                            if plural_index > 0:
                                entry.msgstr_plural[plural_index] = val
                                updated = True
            else:
                # Simple entry
                if entry.msgid in translations:
                    val = translations[entry.msgid]
                    if val and val.strip():
                        entry.msgstr = val
                        updated = True
                        
            # Clean fuzzy flag if successfully updated
            if updated:
                if 'fuzzy' in entry.flags:
                    entry.flags.remove('fuzzy')
                applied_count += 1
                
    print(f"Applied {applied_count} translations to PO entry items.")
    
    # Save the updated PO file
    po.save()
    print(f"Saved updated PO file to {po_path}")
    return True

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