#!/usr/bin/env python

import difflib
import json
from pathlib import Path

import json5


def main():
    script_dir = Path(__file__).parent.resolve()
    # Adjust path relative to the script's location in the aider repo
    litellm_path = script_dir.parent / "../litellm/model_prices_and_context_window.json"
    aider_path = script_dir / "../aider/resources/model-metadata.json"

    if not litellm_path.exists():
        print(f"Error: LiteLLM metadata file not found at {litellm_path}")
        return

    if not aider_path.exists():
        print(f"Error: Aider metadata file not found at {aider_path}")
        return

    try:
        with open(litellm_path, "r") as f:
            litellm_data = json.load(f)
    except json.JSONDecodeError as e:
        print(f"Error decoding JSON from {litellm_path}: {e}")
        return
    except Exception as e:
        print(f"Error reading {litellm_path}: {e}")
        return

    try:
        # Use json5 for the aider metadata file as it might contain comments
        with open(aider_path, "r") as f:
            aider_data = json5.load(f)
    except json.JSONDecodeError as e:
        print(f"Error decoding JSON from {aider_path}: {e}")
        return
    except Exception as e:
        print(f"Error reading {aider_path}: {e}")
        return

    litellm_keys = set(litellm_data.keys())
    aider_keys = set(aider_data.keys())

    keys_to_remove = set()
    common_keys = sorted(list(litellm_keys.intersection(aider_keys)))

    if common_keys:
        print("Comparing common models found in both files:\n")
        for key in common_keys:
            print(f"--- {key} (aider) ---")
            print(f"+++ {key} (litellm) +++")

            litellm_entry = litellm_data.get(key, {})
            aider_entry = aider_data.get(key, {})

            # Convert dicts to formatted JSON strings for comparison
            litellm_json = json.dumps(litellm_entry, indent=4, sort_keys=True).splitlines()
            aider_json = json.dumps(aider_entry, indent=4, sort_keys=True).splitlines()

            # Generate unified diff
            diff = difflib.unified_diff(
                aider_json,
                litellm_json,
                fromfile=f"{key} (aider)",
                tofile=f"{key} (litellm)",
                lineterm="",
                n=max(len(litellm_json), len(aider_json)),  # Show all lines
            )

            # Print the diff, skipping the header lines generated by unified_diff
            diff_lines = list(diff)[2:]
            if not diff_lines:
                print("(No differences found)")
            else:
                for line in diff_lines:
                    # Add color for better readability (optional, requires a library
                    # like 'termcolor' or manual ANSI codes)
                    # Simple +/- indication is standard for diffs
                    print(line)
            print("\n" + "=" * 40)

            # Ask user if they want to remove the entry from aider's metadata
            response = (
                input(f"Remove '{key}' from aider/resources/model-metadata.json? (y/N): ")
                .strip()
                .lower()
            )
            if response == "y":
                keys_to_remove.add(key)
                print(f"Marked '{key}' for removal.")
            else:
                print(f"Keeping '{key}'.")
            print("-" * 40 + "\n")  # Separator for the next model

    else:
        print("No common models found between the two files.")
        return  # Exit if no common keys

    # Remove marked keys after iterating through all common models
    if keys_to_remove:
        print("\nRemoving marked entries from aider data...")
        removed_count = 0
        for key in keys_to_remove:
            if key in aider_data:
                del aider_data[key]
                print(f" - Removed {key}")
                removed_count += 1

        if removed_count > 0:
            # Write the modified data back to the aider metadata file
            try:
                with open(aider_path, "w") as f:
                    # Use json.dump for standard, clean JSON output
                    # Using sort_keys=True for consistent ordering
                    json.dump(aider_data, f, indent=4, sort_keys=True)
                    # Add a trailing newline for POSIX compatibility
                    f.write("\n")
                print(f"\nSuccessfully updated {aider_path} with {removed_count} removal(s).")
            except Exception as e:
                print(f"\nError writing updated data to {aider_path}: {e}")
        else:
            print("\nNo entries were actually removed (perhaps they were already gone?).")
    else:
        print("\nNo entries marked for removal.")


if __name__ == "__main__":
    main()