from urllib.parse import urlparse, parse_qs


def _behavior_to_agent_id(behavior_name: str, unique_id: int) -> str:
    return f"{behavior_name}?agent_id={unique_id}"


def _agent_id_to_behavior(agent_id: str) -> str:
    return agent_id.split("?agent_id=")[0]


def _unwrap_batch_steps(batch_steps, behavior_name):
    decision_batch, termination_batch = batch_steps
    decision_id = [
        _behavior_to_agent_id(behavior_name, i) for i in decision_batch.agent_id
    ]
    termination_id = [
        _behavior_to_agent_id(behavior_name, i) for i in termination_batch.agent_id
    ]
    agents = decision_id + termination_id
    obs = {
        agent_id: [batch_obs[i] for batch_obs in termination_batch.obs]
        for i, agent_id in enumerate(termination_id)
    }
    if decision_batch.action_mask is not None:
        obs.update(
            {
                agent_id: {
                    "observation": [batch_obs[i] for batch_obs in decision_batch.obs],
                    "action_mask": [mask[i] for mask in decision_batch.action_mask],
                }
                for i, agent_id in enumerate(decision_id)
            }
        )
    else:
        obs.update(
            {
                agent_id: [batch_obs[i] for batch_obs in decision_batch.obs]
                for i, agent_id in enumerate(decision_id)
            }
        )
    obs = {k: v if len(v) > 1 else v[0] for k, v in obs.items()}
    dones = {agent_id: True for agent_id in termination_id}
    dones.update({agent_id: False for agent_id in decision_id})
    rewards = {
        agent_id: termination_batch.reward[i]
        for i, agent_id in enumerate(termination_id)
    }
    rewards.update(
        {agent_id: decision_batch.reward[i] for i, agent_id in enumerate(decision_id)}
    )
    cumulative_rewards = {k: v for k, v in rewards.items()}
    infos = {}
    for i, agent_id in enumerate(decision_id):
        infos[agent_id] = {}
        infos[agent_id]["behavior_name"] = behavior_name
        infos[agent_id]["group_id"] = decision_batch.group_id[i]
        infos[agent_id]["group_reward"] = decision_batch.group_reward[i]
    for i, agent_id in enumerate(termination_id):
        infos[agent_id] = {}
        infos[agent_id]["behavior_name"] = behavior_name
        infos[agent_id]["group_id"] = termination_batch.group_id[i]
        infos[agent_id]["group_reward"] = termination_batch.group_reward[i]
        infos[agent_id]["interrupted"] = termination_batch.interrupted[i]
    id_map = {agent_id: i for i, agent_id in enumerate(decision_id)}
    return agents, obs, dones, rewards, cumulative_rewards, infos, id_map


def _parse_behavior(full_behavior):
    parsed = urlparse(full_behavior)
    name = parsed.path
    ids = parse_qs(parsed.query)
    team_id: int = 0
    if "team" in ids:
        team_id = int(ids["team"][0])
    return name, team_id