#!/usr/bin/env python3
import argparse
import json
import os
import sys
import requests

API_URL = "https://habitica.com/api/v3/user"

def get_habitica_headers():
    user_id = os.getenv("HABITICA_USER_ID")
    api_token = os.getenv("HABITICA_API_TOKEN")

    if not user_id or not api_token:
        print("Error: Environment variables HABITICA_USER_ID and HABITICA_API_TOKEN must be set.", file=sys.stderr)
        sys.exit(1)

    return {
        "x-api-user": user_id,
        "x-api-key": api_token,
        "x-client": f"{user_id}-CustomizationScript"
    }
    
def fetch_habitica_gear():
    """Fetches only the costume, background, pet, and mount from the user profile."""
    try:
        response = requests.get(API_URL, headers=get_habitica_headers())
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        print(f"Error communicating with Habitica API: {e}", file=sys.stderr)
        sys.exit(1)

    data = response.json().get("data", {})

    # Extract specific fields
    items = data.get("items", {})
    gear = items.get("gear", {})
    costume = gear.get("costume", {})
    background = data.get("preferences", {}).get("background", "default")
    pet = items.get("currentPet", "")
    mount = items.get("currentMount", "")

    customizations = {
        "costume": costume,
        "background": background,
        "pet": pet,
        "mount": mount
    }

    print(customizations)
    return customizations

def save_costume(filename):
    # print("Fetching equipment and costume data from Habitica...")
    gear_data = fetch_habitica_gear()
    with open(filename, "w", encoding="utf-8") as f:
        json.dump(gear_data, f, indent=4)


def equip(type, name):
    try:
        response = requests.post(f"{API_URL}/equip/{type}/{name}", 
                                 headers=get_habitica_headers())
        response.raise_for_status()
        print(f"Successfully equiped '{name}'.")
    except requests.exceptions.RequestException as e:
        print(f"Error communicating with Habitica API: {e}", file=sys.stderr)
        if hasattr(e, 'response') and e.response is not None:
            print(f"Server response: {e.response.text}", file=sys.stderr)
        sys.exit(1)

def set_background(name):
    bg_payload = {
        "preferences.background": name
    }
    try:
        res = requests.put(f"{API_URL}", 
                           headers=get_habitica_headers(), json=bg_payload)
        res.raise_for_status()
        print(f"Successfully set background to: {name}")
    except requests.exceptions.RequestException as e:
        print(f"Error setting background: {e}", file=sys.stderr)
        if hasattr(e, 'response') and e.response is not None:
            print(f"Server response: {e.response.text}", file=sys.stderr)

def set_costume(data):
    for k,v in data.get("costume").items():
        if v and not v.endswith("base_0"):
            equip("costume", v)

    if "pet" in data:
        equip("pet", data.get("pet"))

    # if "mount" in data:
    #     equip("mount", data.get("mount"))
        
    # TODO: Not sure how to set the background!
    if "background" in data:
        set_background(data["background"])

def load_costume(filename):
    if os.path.exists(filename):
        with open(filename, "r", encoding="utf-8") as f:
            gear_data = json.load(f)
            # Update defaults with values from the file
            set_costume(gear_data)

def store_all(filename):
    try:
        response = requests.get(API_URL, headers=get_habitica_headers())
        response.raise_for_status()
    except requests.exceptions.RequestException as e:
        print(f"Error communicating with Habitica API: {e}", file=sys.stderr)
        sys.exit(1)

    data = response.json().get("data", {})
    with open(filename, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)

def parse_arguments():
    parser = argparse.ArgumentParser(description="Costume configuration script")
    parser.add_argument(
        "--load",
        type=str,
        help="Path to a JSON file containing saved costume values"
    )
    parser.add_argument(
        "--save",
        type=str,
        help="Path to a JSON file to save your costume values"
    )
    parser.add_argument(
        "-a", "--all",
        action=argparse.BooleanOptionalAction,
        help="Download and store all data in 'all.json'"
    )
    return parser.parse_args()

def main():
    args = parse_arguments()
    
    if args.save:
        save_costume(args.save)
    elif args.load:
        load_costume(args.load)
    elif args.all:
        store_all("all.json")
    
if __name__ == "__main__":
    main()