import os
import json
import pandas as pd

from uot.chat_utils import import_prompts_by_task
from uot.uot import UoTNode


class TBTask:
    def __init__(self, args):
        self.__dict__.update(vars(args))
        self.free_answer = False
        self.max_turn = 20
        self.prompts = import_prompts_by_task("tb")
        self.set = []
        self.data = self.load_dataset(args.dataset)
        self.root = None

    def load_dataset(self, name):
        if name == "FloDial":
            trouble_name = {'car_electrical_failure|chart-9': 'Parasitic drain', 'car_electrical_failure|chart-10': 'Battery', 'car_electrical_failure|chart-4': 'Blown fuse', 'car_electrical_failure|chart-5': 'Voltage issue', 'car_electrical_failure|chart-23': 'Corroded battery connectors', 'car_electrical_failure|chart-12': 'Voltage regulator', 'car_electrical_failure|chart-24': 'Open circuit', 'car_electrical_failure|chart-15': 'Bulb failure', 'car_electrical_failure|chart-17': 'bulb replacement', 'car_electrical_failure|chart-18': 'multiple light failure', 'lcd_problem|chart-20': 'Video problem', 'lcd_problem|chart-15': 'LCD fault', 'lcd_problem|chart-12': 'Backlight', 'lcd_problem|chart-24': 'Overheating', 'lcd_problem|chart-28': 'Resolution', 'lcd_problem|chart-21': 'Loose display connectors', 'lcd_problem|chart-29': 'Wiring issue', 'lcd_problem|chart-13': 'Inverter', 'lcd_problem|chart-26': 'LCD leakage', 'lcd_problem|chart-17': 'Screen rotation', 'power|chart-22': 'Malfunctioning internal device', 'power|chart-8': 'LCD display issue', 'power|chart-4': 'Beep Codes', 'power|chart-20': 'Battery issue', 'power|chart-11': 'Power button issue', 'power|chart-13': 'Motherboard failure', 'power|chart-23': 'Motherboard failure', 'power|chart-6': 'External device interference', 'power|chart-15': 'Cord issue', 'power|chart-12': 'Loose connections', 'power|chart-24': 'Power connector issue', 'wireless|chart-27': 'ISP issue', 'wireless|chart-3': 'Wireless adapter issue', 'wireless|chart-16': 'ISP problem', 'wireless|chart-8': 'Reset', 'wireless|chart-29': 'Router reset', 'wireless|chart-28': 'Interference', 'wireless|chart-18': 'Network connectivity issue', 'wireless|chart-9': 'Network visibility', 'wireless|chart-25': 'Incompatibility', 'wireless|chart-23': 'Driver issue', 'laptop_overheating|chart-18': 'Overheating shutdown', 'laptop_overheating|chart-20': "The issue's name is not specified in the given statement", 'laptop_overheating|chart-15': 'Power management', 'laptop_overheating|chart-17': 'Overheating', 'laptop_overheating|chart-19': 'Hardware failure', 'laptop_overheating|chart-27': 'Overheating', 'laptop_overheating|chart-4': 'Battery issue', 'laptop_overheating|chart-23': 'Burnt motherboard', 'laptop_overheating|chart-26': 'Overheating', 'laptop_drive|chart-12': 'Hard drive failure', 'laptop_drive|chart-27': 'LCD malfunction', 'laptop_drive|chart-21': 'Controller failure', 'laptop_drive|chart-23': 'Virus/malware identification and OS restoration', 'laptop_drive|chart-11': 'Loose connector', 'laptop_drive|chart-18': 'Unknown', 'laptop_drive|chart-5': 'Hard drive failure', 'laptop_drive|chart-28': 'Battery issue', 'laptop_drive|chart-22': 'Windows installation problem', 'laptop_drive|chart-15': 'Boot order issue', 'ticking|chart-25': 'brake adjusters', 'ticking|chart-14': 'CV joint failure', 'ticking|chart-11': 'Expansion/contraction', 'ticking|chart-16': 'Loose lug nuts', 'ticking|chart-21': 'Warped rotor', 'ticking|chart-22': 'Debris', 'ticking|chart-4': 'Faulty ball joint', 'ticking|chart-9': 'Exhaust leak', 'engine_overheats|chart-20': 'Antifreeze leak', 'engine_overheats|chart-31': 'Leak', 'engine_overheats|chart-35': 'Coolant temperature light issue', 'engine_overheats|chart-30': 'Water pump failure', 'engine_overheats|chart-16': 'Overheating', 'engine_overheats|chart-8': 'fan or temperature sensor failure', 'engine_overheats|chart-10': 'Clogged radiator', 'engine_overheats|chart-29': 'Water pump leak', 'engine_overheats|chart-34': 'Pressure relief system failure', 'engine_overheats|chart-12': 'engine flush', 'engine_overheats|chart-19': 'Thermostat', 'engine_overheats|chart-6': 'overheating', 'car_wont_start|chart-33': "The issue's name is a broken ignition switch", 'car_wont_start|chart-13': 'Short circuit', 'car_wont_start|chart-17': 'cap erosion', 'car_wont_start|chart-22': 'Fuel injector problem', 'car_wont_start|chart-21': 'Carburetor issue', 'car_wont_start|chart-30': 'Choke stuck', 'car_wont_start|chart-7': 'Starter relay failure', 'car_wont_start|chart-31': 'sensor', 'car_wont_start|chart-16': 'Fuel distribution', 'car_wont_start|chart-19': 'Fuel issue', 'car_wont_start|chart-4': 'Dead battery', 'brake_problem|chart-6': 'ABS failure', 'brake_problem|chart-29': 'Brake wear', 'brake_problem|chart-33': 'Worn parts', 'brake_problem|chart-7': "The issue's name is not specified in the given statement", 'brake_problem|chart-27': "The issue's name is brake pad vibration", 'brake_problem|chart-16': 'brake problem', 'brake_problem|chart-24': 'High spot on the rotor', 'brake_problem|chart-13': 'rust or malfunctioning return spring', 'brake_problem|chart-34': 'ABS malfunction', 'brake_problem|chart-37': 'Wheel bearing', 'brake_problem|chart-11': 'Stuck piston or over-adjusted shoes', 'brake_problem|chart-25': 'Spring clips', 'brake_problem|chart-36': 'Air in brake line system', 'car_electrical_failure|chart-27': 'Short circuit', 'car_electrical_failure|chart-13': 'Drive belt', 'lcd_problem|chart-5': 'Video processor', 'lcd_problem|chart-23': "The issue's name is CCFL malfunction", 'lcd_problem|chart-4': 'Power failure', 'power|chart-28': 'AC adapter malfunction', 'power|chart-16': 'DC power jack failure', 'wireless|chart-17': 'Router malfunction', 'wireless|chart-10': 'Connectivity', 'wireless|chart-7': 'Driver issue', 'laptop_overheating|chart-10': 'Overheating', 'laptop_overheating|chart-9': 'BIOS upgrade causing overheating', 'laptop_drive|chart-30': 'Motherboard malfunction', 'laptop_drive|chart-8': 'Hardware failure or corrupted driver', 'laptop_drive|chart-33': 'corrupted Windows installation', 'ticking|chart-26': 'Transmission issue', 'ticking|chart-23': 'Loose wire', 'engine_overheats|chart-17': 'Overheating', 'engine_overheats|chart-33': 'cooling system problem', 'car_wont_start|chart-14': 'Ignition system wiring or voltage regulator issue', 'car_wont_start|chart-8': 'Starter problem', 'car_wont_start|chart-32': 'cracked coil or cracked distributor', 'brake_problem|chart-26': 'Loose bolt', 'brake_problem|chart-4': 'Brake fluid overflow', 'brake_problem|chart-8': 'jamming', 'brake_problem|chart-32': 'Emergency brake switch', 'car_electrical_failure|chart-21': 'Open circuit', 'car_electrical_failure|chart-26': 'Blown fusible link', 'lcd_problem|chart-9': 'Hard drive issue', 'lcd_problem|chart-8': 'Display issue', 'power|chart-29': 'Non-functional power outlet', 'power|chart-27': 'connector malfunction', 'wireless|chart-12': 'Password', 'wireless|chart-24': 'Malfunction', 'laptop_battery|chart-18': 'Overdischarge', 'laptop_battery|chart-14': 'battery memory effect', 'laptop_battery|chart-12': 'faulty battery', 'laptop_overheating|chart-5': 'Overheating', 'laptop_overheating|chart-22': 'malfunctioning motherboard or RAM', 'car_steering|chart-19': 'Bent part', 'car_steering|chart-37': 'Bent axle', 'car_steering|chart-22': 'Alignment', 'car_steering|chart-33': 'Wheel imbalance', 'laptop_drive|chart-25': 'Corrupted hard drive', 'laptop_drive|chart-13': 'Battery issue', 'ticking|chart-3': "The issue's name is not provided in the given statement", 'ticking|chart-20': 'Pebbles in wheel covers', 'engine_overheats|chart-27': 'Heater core malfunction', 'engine_overheats|chart-32': 'Radiator leak', 'engine_overheats|chart-14': 'Thermostat', 'car_wont_start|chart-24': 'spark control system', 'car_wont_start|chart-26': 'Blink codes issue', 'car_wont_start|chart-6': 'Oxidation', 'brake_problem|chart-15': 'Cable issue', 'brake_problem|chart-22': 'Brake pad issue', 'laptop_battery|chart-4': 'power adapter', 'laptop_battery|chart-20': 'Power drain', 'laptop_battery|chart-22': 'Battery issue', 'laptop_battery|chart-23': 'Battery issue', 'laptop_battery|chart-16': 'battery charger problem', 'laptop_battery|chart-24': 'Power plan', 'laptop_battery|chart-10': 'Battery replacement', 'car_steering|chart-11': 'Worn steering coupling', 'car_steering|chart-8': 'worn lower ball joint', 'car_steering|chart-16': 'Strut failure', 'car_steering|chart-21': 'Alignment', 'car_steering|chart-36': 'CV joints', 'car_steering|chart-28': 'Leak', 'car_steering|chart-30': 'Mechanical', 'car_steering|chart-34': 'Tire damage', 'car_steering|chart-3': 'Loose lug nuts', 'car_steering|chart-13': 'Power steering pump issue', 'car_steering|chart-14': 'Loose rack', 'car_steering|chart-5': 'Cotter pin issue', 'car_steering|chart-25': 'Loose belt', 'laptop_battery|chart-6': 'BIOS power controller', 'laptop_battery|chart-7': 'Battery drain', 'car_steering|chart-26': 'Pump failure', 'car_steering|chart-15': 'Tie rod'}
            self.set = list(trouble_name.values()) if self.open_set_size <= 0 else self.set
        else:
            raise NotImplementedError
        return json.loads(os.path.join(os.path.dirname(__file__), f"../data/{name}.json").read())

    def create_root(self, root=None):
        if not root:
            self.root = UoTNode("ROOT", True, self.set, None, self.guesser_model)
        else:
            root.set_config(self.n_extend_layers, not self.none_acc_reward, self.expected_reward_method)
            self.root = root


def load_flodial_dataset(file_path):
    repo_dataset = []
    df = pd.read_excel(file_path)
    for _, row in df.iterrows():
        repo_dataset.append({'self_repo': row['Self-repo'], 'target': row['Item']})
    return repo_dataset
