import argparse
import json
import os
import random
import select
import signal
import subprocess
import sys
import time
import traceback
from pathlib import Path

import numpy as np
import requests
import torch
import yaml

from constants.constants import LM_CONFIGS, URLS_DIR
from utils.file_utils import get_attribute_from_dict


# ===============================================================================
# TGI deployment aux functions #TODO move or delete this
# ===============================================================================
def deploy_tgi(model_path, quantize, num_shard):
    print(f"\n\nDeploying TGI locally for model {model_path}\n")
    process = subprocess.Popen(
        [
            "text-generation-launcher",
            "--model-id",
            model_path,
            "--quantize",
            quantize,
            "--num-shard",
            str(num_shard),
            "--port8080--master-port8080--master-addrlocalhost",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )

    print("TGI has been started and is running in the background.")
    print(f"Process ID: {process.pid}")
    return process


def terminate_process(process):
    if process and process.poll() is None:
        print(f"Terminating TGI process...{process.pid}")
        os.kill(process.pid, signal.SIGTERM)
        process.wait()
        print("Process terminated")


def signal_handler(process, sig, frame):
    print(f"SIGINT received, terminating the process...{process.pid}")
    if process and process.poll() is None:  # Check if the process is still running
        os.kill(process.pid, signal.SIGTERM)
        process.wait()
        print("Process terminated gracefully.")
    sys.exit(0)


# ===============================================================================
# Agent and Model config helpers
# ===============================================================================


def load_agent_config(agent_config_file: str) -> dict:
    with open(agent_config_file, "r") as file:
        agents_configs = yaml.safe_load(file)
    return agents_configs


def get_lm_config(model, gen_config_alias, lm_config_file: str = LM_CONFIGS) -> dict:
    with open(lm_config_file, "r") as file:
        all_lm_configs = yaml.safe_load(file)
    return all_lm_configs[model][gen_config_alias]


def resolve_inheritance(agents_configs: dict):
    # For each agent
    for agent_str, agent_config in agents_configs.items():
        # For ach attribute
        for key, subconfig in agent_config.items():
            if not isinstance(subconfig, dict):
                continue

            if "_inherit_from" in subconfig:
                inherit_from = subconfig["_inherit_from"]
                agents_configs[agent_str][key] = agents_configs[inherit_from][key].copy()
    return agents_configs


def get_agent_attribute(agent_config_path: str, attribute: str):
    agents_configs = load_agent_config(agent_config_path)
    return get_attribute_from_dict(attribute, agents_configs)


def get_agent_config(agent_config_path: str, active_modules_str: str = ""):
    # Load general agent configurations
    agents_configs = load_agent_config(agent_config_path)

    active_modules = []
    if active_modules_str:
        active_modules = active_modules_str.lower().split(":")

    # Set model and LM Config for each agent
    for module, config in agents_configs.items():
        if active_modules and not any(m in module.lower() for m in active_modules):
            continue

        # If not LLM-based agent, skip
        if "lm_config" not in config:
            continue

        # If LM configs inherited from other modules, inherit
        if "_inherit_from" in config["lm_config"]:
            inherit_from = config["lm_config"]["_inherit_from"]
            config["lm_config"] = agents_configs[inherit_from]["lm_config"]

        # Fill LM config with actual values
        config["lm_config"].update(
            get_lm_config(
                model=config["lm_config"]["model"],
                gen_config_alias=config["lm_config"]["gen_config_alias"],
            )
        )
    return agents_configs


def dump_config(args: argparse.Namespace, logger) -> None:
    config_file = Path(args.result_dir) / "config.json"
    if not config_file.exists():
        with open(config_file, "w") as f:
            json.dump(vars(args), f, indent=4)
            logger.info(f"Dump config to {config_file}")


# ===============================================================================
# Other helpers
# ===============================================================================


def set_seed(seed: int | None = None) -> None:
    if seed is None:
        return
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
