import os
import subprocess
import json
import requests

API_KEY = os.environ.get("ARK_API_KEY")
url = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
API_KEY = os.environ.get("ARK_API_KEY")
API_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
room = {"Kitchen", "Living Room", "Bedroom", "Bathroom"}
room_id = {"1", "2", "3", "4"}
action_list = {
    'OpenObject',
    'CloseObject',
    'PickupObject',
    'PutObject',
    'CleanObject',
    'HeatObject',
    'CoolObject',
    'ToggleObjectOn',
    'ToggleObjectOff',
    'SliceObject'}
RECEPTACLES = {
    'BathtubBasin',
    'Bowl',
    'Cup',
    'Drawer',
    'Mug',
    'Plate',
    'Shelf',
    'SinkBasin',
    'Box',
    'Cabinet',
    'CoffeeMachine',
    'CounterTop',
    'Fridge',
    'GarbageCan',
    'HandTowelHolder',
    'Microwave',
    'Pot',
    'StoveBurner',
    'SideTable',
    'ToiletPaperHanger',
    'TowelHolder',
    'Safe',
    'ArmChair',
    'Toilet',
    'Dresser',
    'LaundryHamper',
    'Toaster',
}
Openable = {
    'Box',
    'Cabinet',
    'Drawer',
    'Fridge',
    'Microwave',
    'Safe',
}
objects = {
    'AlarmClock', 'Apple', 'AppleSliced', 'BaseballBat', 'BasketBall', 'Book', 'Bowl', 'Box', 'Bread', 'BreadSliced',
    'ButterKnife', 'CD', 'Candle', 'CellPhone', 'Cloth', 'CreditCard', 'Cup', 'DeskLamp', 'DishSponge', 'Egg', 'Faucet',
    'FloorLamp', 'Fork', 'Glassbottle', 'HandTowel', 'HousePlant', 'Kettle', 'KeyChain', 'Knife', 'Ladle', 'Laptop',
    'LaundryHamperLid', 'Lettuce', 'LettuceSliced', 'LightSwitch', 'Mug', 'Newspaper', 'Pan', 'PaperTowel',
    'PaperTowelRoll', 'Pen', 'Pencil', 'PepperShaker', 'Pillow', 'Plate', 'Plunger', 'Pot', 'Potato', 'PotatoSliced',
    'RemoteControl', 'SaltShaker', 'ScrubBrush', 'ShowerDoor', 'SoapBar', 'SoapBottle', 'Spatula', 'Spoon',
    'SprayBottle', 'Statue', 'StoveKnob', 'TeddyBear', 'Television', 'TennisRacket', 'TissueBox', 'ToiletPaper',
    'ToiletPaperRoll', 'Tomato', 'TomatoSliced', 'Towel', 'Vase', 'Watch', 'WateringCan', 'WineBottle'
}
map_save_large_objects = [
    'ArmChair', 'BathtubBasin', 'Bed', 'Cabinet', 'Cart', 'CoffeeMachine', 'CoffeeTable', 'CounterTop',
    'Desk', 'DiningTable', 'Drawer', 'Dresser', 'Fridge', 'GarbageCan', 'Microwave', 'Ottoman', 'Safe',
    'Shelf', 'SideTable', 'SinkBasin', 'Sofa', 'StoveBurner', 'TVStand', 'Toilet'
]
# 将 receptacles 中不在 objects 里的元素加进去
merged_objects = sorted(objects.union(set(map_save_large_objects)))
rooms = {
    "Kitchen": {
        "ALL": [
            "Apple", "AppleSliced", "Bowl", "Bread", "BreadSliced", "ButterKnife",
            "Cabinet", "CoffeeMachine", "CounterTop", "Egg", "Fork", "Fridge", "Kettle", "Knife", "Lettuce",
            "LettuceSliced",
            "LightSwitch", "Microwave", "Pan", "PepperShaker", "Plate", "Pot",
            "Potato", "PotatoSliced", "SaltShaker", "Sink", "Spatula", "Spoon",
            "StoveBurner", "StoveKnob", "Toaster", "Tomato", "TomatoSliced", "WineBottle"
        ],
        "some": [
            "Box", "Chair", "Candle", "Cup", "Drawer", "DiningTable", "GarbageCan", "Glassbottle", "HousePlant",
            "KeyChain", "Laptop", "Mug", "Newspaper", "Painting", "PaperTowel",
            "PaperTowelRoll", "Pen", "Pencil", "Safe", "ScrubBrush", "SoapBottle",
            "SprayBottle", "Statue", "Towel", "Vase", "Watch", "WateringCan", "Window"
        ]
    },
    "Living Room": {
        "ALL": [
            "ArmChair", "Chair", "FloorLamp", "LightSwitch", "RemoteControl",
            "SideTable", "Sofa", "Television"
        ],
        "some": [
            "BaseballBat", "BasketBall", "Box", "Candle", "Cup", "Desk", "DeskLamp",
            "GarbageCan", "Glassbottle", "HousePlant", "KeyChain", "Laptop", "Mug",
            "Newspaper", "Ottoman", "Painting", "PaperTowel", "PaperTowelRoll", "Pen",
            "Pencil", "Pillow", "Safe", "Statue", "TeddyBear", "TennisRacket",
            "TissueBox", "Vase", "Watch", "WateringCan", "Window", "WineBottle"
        ]
    },
    "Bedroom": {
        "ALL": [
            "AlarmClock", "Bed", "Desk", "DeskLamp", "Dresser", "LightSwitch",
            "Laptop", "PaperTowel", "PaperTowelRoll", "Pen", "Pencil", "Pillow", "TeddyBear"
        ],
        "some": [
            "ArmChair", "BaseballBat", "BasketBall", "Box", "Candle", "Chair",
            "FloorLamp", "GarbageCan", "Glassbottle", "HousePlant", "KeyChain",
            "LaundryHamper", "LaundryHamperLid", "Mug", "Newspaper", "Ottoman",
            "Painting", "Poster", "RemoteControl", "Safe", "Statue", "Television",
            "TennisRacket", "TissueBox", "Towel", "Vase", "Watch", "WateringCan", "Window"
        ]
    },
    "Bathroom": {
        "ALL": [
            "BathtubBasin", "Cabinet", "Drawer", "Faucet", "HandTowel", "HandTowelHolder",
            "LightSwitch", "Plunger", "ShowerCurtain", "ShowerDoor", "ShowerGlass",
            "ShowerHead", "Sink", "SoapBar", "SoapBottle", "Toilet", "ToiletPaper",
            "ToiletPaperHanger", "Towel", "TowelHolder"
        ],
        "some": [
            "Box", "Candle", "CounterTop", "Cup", "GarbageCan", "Glassbottle",
            "LaundryHamper", "LaundryHamperLid", "Mug", "Newspaper", "Painting",
            "PaperTowel", "PaperTowelRoll", "Pen", "Pencil", "ScrubBrush", "SprayBottle",
            "Statue", "TissueBox", "Vase", "Watch", "Window"
        ]
    }
}
'''def call_ark_from_py36(user_input):
    payload = {
        "api_key": "6d67df38-2dde-4e4e-951a-22e1cd117ef1",  # Or use os.environ.get("ARK_API_KEY")
        "model": "deepseek-r1-250120",
        "messages": [
            {"role": "user", "content": user_input}
        ]
    }

    python38_path = os.path.expanduser("/scratch/lijh/.conda/envs/janus/bin/python")
    ark_runner_script = "ark_example.py"

    process = subprocess.Popen(
        [python38_path, ark_runner_script],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

    out, err = process.communicate(input=json.dumps(payload).encode("utf-8"))

    if err and err.strip():
        print("❌ Subprocess Error:", err.decode())
    else:
        result = json.loads(out.decode("utf-8"))
        print("✅ Ark Reply:", result)
        return result['result'].strip()'''


def call_ark_from_py36(user_input):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    data = {
        "model": "doubao-1-5-pro-32k-250115",
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": user_input}
        ],
        "temperature": 0.1
    }
    response = requests.post(API_URL, headers=headers, data=json.dumps(data))
    if response.status_code == 200:
        result = response.json()
        return result['choices'][0]['message']['content'].strip()
    else:
        print("调用失败：", response.status_code)
        print(response.text)
        return ""


def location_predict(goal_long, goal_short, list_of_actions_pointer_s, list_of_actions_s, detected_list):
    template = """
        You are an intelligent agent operating in a home-like virtual environment.

Here are the known objects/locations in the environment:
- smallObjects: {Objects1}
- largeObjects: {Objects2}

Your current task is: {Task}
Your current task description is: {TaskDescription}
You are currently following the plan: {InitialPlan}
You are now stuck at step "{pointer}{Stuck}" .

You need to determine the most likely location of the object: "{object_prediction}"

Instructions:
- Your answer **must** be one of the entries from the **largeObjects** list provided above, or "None".
- Only choose a location if it is clearly implied or explicitly stated in the task or task description.
- If there is **no indication** of the object's location, return **None**.
- You **must not** include any explanation or extra information in your response.
- Your final output should be a **single word or phrase** (i.e."SideTable" or "None").


Now, based on the task description, where can the object "{object_prediction}" most likely be found?


"""

    prompt = template.format(
        Objects1=objects,
        Objects2=map_save_large_objects,
        Openable=Openable,
        Task=goal_short,
        Objects=merged_objects,
        TaskDescription=goal_long,
        InitialPlan=list_of_actions_s,
        pointer=list_of_actions_pointer_s,
        Stuck=list_of_actions_s[list_of_actions_pointer_s],
        ObservedObjects=detected_list,
        ActionList=action_list,
        object_prediction=list_of_actions_s[list_of_actions_pointer_s][0]
    )
    print("list_of_actions_s[list_of_actions_pointer_s][0]", list_of_actions_s[list_of_actions_pointer_s][0])
    print(list_of_actions_s[list_of_actions_pointer_s])

    # print("call_ark_from_py36")
    # print(prompt)
    # print(result)
    # result=call_ark_from_py36(prompt)
    result = call_ark_from_py36(prompt)
    print(result)
    # clean_plan= extract_plan_from_reply(result)
    clean_plan = postprocess_location(result)
    return clean_plan


import re


def postprocess_location(result):
    """
    处理 replan 返回的字符串 result：
    - 去除多余空格和标点
    - 判断是否是 largeObjects 中的元素
    """
    cleaned = result.strip().strip(" .,:;!?()[]\"'")  # 去除前后标点和空格
    for obj in map_save_large_objects:
        if cleaned.lower() == obj.lower():
            return obj  # 保留原始拼写
    return "None"


def main():
    goal_long = "Turn around and walk to the right around the living room table and turn left and look up to face the statue on the shelf above the couch.', 'Pick up the statue on the shelf above the couch.', 'Turn around to face the living room table.', 'Place the statue on the right side of the living room table.', 'Turn right and walk across the room to face the table to the left of the chair.', 'Pick up the statue on the table.', 'Turn around and walk around to face the living room table.', 'Place the statue beside the first statue on the living room table"
    goal_short = "To move two statues to the living room table"
    list_of_actions_pointer_s = 2
    list_of_actions_s = [('Statue', 'PickupObject'), ('CoffeeTable', 'PutObject'), ('Statue', 'PickupObject'),
                         ('CoffeeTable', 'PutObject')]
    detected_list = ['Fridge', 'CoffeeTable', 'DiningTable']
    step1 = 5
    step2 = 10

    plan = location_predict(goal_long, goal_short, list_of_actions_pointer_s, list_of_actions_s, detected_list, step1,
                            step2)
    print(plan)


if __name__ == "__main__":
    main()