import json
import xml.etree.ElementTree as ET
from elevenlabslib.helpers import *
from elevenlabslib import *
from sys import exit
import shutil
import os
from pydub import AudioSegment
import time

class Voice:
    def __init__(self,r:int,g:int,b:int):
        self.r = r
        self.g = g
        self.b = b
    
    def is_SameVoice(self,r:int,g:int,b:int) -> bool:
        if self.r != r:
            return False
        if self.g != g:
            return False
        if self.b != b:
            return False
        return True
    def GetVoiceNumber(self) -> int:
        if self.is_SameVoice(255,192,0):
            return 1
        elif self.is_SameVoice(0,176,80):
            return 2
        elif self.is_SameVoice(128,128,255):
            return 3
        elif self.is_SameVoice(255,0,0):
            return 4
        elif self.is_SameVoice(0,0,0):
            return 5
        elif self.is_SameVoice(0,176,240):
            return 6
        elif self.is_SameVoice(255,0,255):
            return 7
        elif self.is_SameVoice(255,255,0):
            return 8
        elif self.is_SameVoice(40,40,40):
            return 9
        elif self.is_SameVoice(255,204,0):
            return 10
        else:
            return 11
    def GetVoiceColor(number : int) -> str:
        if number == 1:
            return '\033[33m' #orange
        elif number == 2:
            return '\033[32m' #green
        elif number == 3:
            return '\033[35m' #purple
        elif number == 4:
            return '\033[31m' #red
        elif number == 5:
            return '\033[30m' #black
        elif number == 6:
            return '\033[36m' #Cyan
        elif number == 7:
            return '\033[95m' #pink
        elif number == 8:
            return '\033[93m' #yellow
        elif number == 9:
            return '\033[90m' #darkgrey
        elif number == 10:
            return '\033[91m' #lightred
        else:
            return '\033[94m' #lightblue

def WriteBroadcastTranscript(Broadcast : ET.Element) -> int:

    transcript = open("./automation/tmp/transcript.txt", "w")

    Lines = []

    for LineEntry in Broadcast:
        Line = LineEntry.text
        #if Line has a time tag, remove it
        if '${' in Line[0:8]:
            Line = Line[9:] #remove ${t:x.xx}
        i = Voice(int(LineEntry.attrib['r']),int(LineEntry.attrib['g']),int(LineEntry.attrib['b'])).GetVoiceNumber()
        
        reset = '\033[0m'

        FinalLine = "{Voice" + str(i) +"}" + Line
        print(Voice.GetVoiceColor(i) + FinalLine + reset)
        Lines.append(FinalLine + "\n")
    transcript.writelines(Lines)
    transcript.close
    return 11
        
# def WriteVoiceConfig(N_Voices : int,BroadcastID: str):

#     #config = open("./automation/tmp/config.txt", "a")

#     keys = ["BroadcastID","API_Key","Voices"]
#     values = [BroadcastID]

#     print("Please write your Eleven Labs API Key")
#     Key = input()
#     values.append(Key)

#     Voices = []

#     for i in range(1, N_Voices+1):
#         print("Please write voice" + str(i) + " ID. If voice is not used, leave blank")
#         VoiceID = input()
#         Voices.append(VoiceID)
#     values.append(Voices)
#     ConfigDict = dict(zip(keys,values))

#     with open("./automation/tmp/config.txt", "w") as out:
#         json.dump(ConfigDict,out,indent="")

def GetTextScript(str: str) ->str:
    LastCommand = str.rfind("}") + 1
    return str[LastCommand:]

def GetCharCost()->int:
    transcript = open("./automation/tmp/transcript.txt", "r")

    lines = transcript.readlines()

    charCost=0

    for line in lines:
        if "Voice" in line:
            #print(GetTextScript(line))
            charCost = charCost + len(GetTextScript(line[:-1]))
    
    transcript.close()

    return charCost

def AskContinue():
    print("please write continue")
    reply = input()
    reply = reply.upper()
    if 'CONTINUE' not in reply:
        print("Quitting")
        exit()
    print("Continuing")

def GetBetweenBrackets(str: str) -> list:
    substrings = []
    in_brackets = False
    current_substring = ""
 
    for c in str:
        if c == "{":
            in_brackets = True
        elif c == "}" and in_brackets:
            substrings.append(current_substring)
            current_substring = ""
            in_brackets = False
        elif in_brackets:
            current_substring += c
 
    if current_substring:
        substrings.append(current_substring)
    return substrings

def GetVoiceLines():
    #config = open("./automation/tmp/config.txt")
    config = open("./automation/config.txt")
    
    data = json.load(config)

    user = ElevenLabsUser(data['API_Key'])
    
    voices = []

    for Voice in data['Voices']:
        voices.append(user.get_voices_by_name(Voice['name'])[0])

    transcript = open("./automation/tmp/transcript.txt", "r")

    lines = transcript.readlines()

    SegmentTimes = []

    FinalAudioFile :AudioSegment

    i=0
    for line in lines:
        commands=GetBetweenBrackets(line)
        j=0
        NoSound = True
        for command in commands:
            if "Voice" in command:
                NoSound = False
                TextScript = GetTextScript(line)
                voiceNumber = int(command[5:])
                try:
                    mp3Data = voices[voiceNumber-1].generate_audio_bytes(TextScript[:-1],data['Voices'][voiceNumber-1]['stability'],data['Voices'][voiceNumber-1]['similarity'])
                except:
                    time.sleep(5) #wait a considerable time for 429 to 'Kindly' go away
                    mp3Data = voices[voiceNumber-1].generate_audio_bytes(TextScript[:-1],data['Voices'][voiceNumber-1]['stability'],data['Voices'][voiceNumber-1]['similarity']) #retry, if it fails, f*ck it
                dst = "./automation/tmp/lines/" + str(i) + "_" + str(j) + ".ogg"
                save_bytes_to_path(dst,mp3Data)
            else:
                #Copy command from automation/SoundEffects to tmp/lines as i_j.ogg
                NoSound = False
                src = "./automation/SoundEffects/" + command
                dst = "./automation/tmp/lines/" + str(i) + "_" + str(j) + ".ogg"
                shutil.copyfile(src,dst)
            j = j+1
        #concatenate audiofiles
        sound : AudioSegment
        if NoSound:
            sound = AudioSegment.silent(duration=3000)
        else:
            for clip in os.listdir("./automation/tmp/lines"):
                if clip.startswith(str(i)+"_"):
                    path = "./automation/tmp/lines/"+clip
                    if clip == (str(i)+"_0.ogg"):
                        sound = AudioSegment.from_ogg(path)
                    else:
                        sound1 = AudioSegment.from_ogg(path)
                        sound = AudioSegment.append(sound,sound1,0)
        sound.export(out_f=("./automation/tmp/lines/"+str(i)+".ogg"),format='ogg')
        SegmentTimes.append(sound.duration_seconds)
        
        if i == 0:
            FinalAudioFile = AudioSegment.from_ogg(("./automation/tmp/lines/"+str(i)+".ogg"))
        else:
            sound1 = AudioSegment.from_ogg(("./automation/tmp/lines/"+str(i)+".ogg"))
            FinalAudioFile = FinalAudioFile.append(sound1,0)

        i = i+1
    FinalAudioFile.export(out_f=("./media/sound/Broadcasts/"+data['BroadcastID']+".ogg"),format='ogg')
    transcript.close()
    config.close()
    return SegmentTimes

def AddSoundIDToScript():
    config = open("./automation/config.txt")
    data = json.load(config)

    id = data['BroadcastID']
    #is ID already in RadioWavs?
    IDinRadioWavs = False
    with open('./media/scripts/sounds_RadioWavs.txt') as f:
        if id in f.read():
            IDinRadioWavs = True
    if not IDinRadioWavs:
        #Remove last char (closing "}")
        with open('./media/scripts/sounds_RadioWavs.txt', 'rb+') as fh:
            fh.seek(-1, 2)
            fh.truncate()
        with open('./media/scripts/sounds_RadioWavs.txt', 'a') as f:
            f.write(f'\n\tsound {id}\n\t{{\n\t\tcategory = Object,\n\t\tclip\n\t\t{{\n\t\t\tfile = media/sound/Broadcasts/{id}.ogg,\n\t\t\tdistanceMin = 5,\n\t\t\tdistanceMax = 50,\n\t\t\treverbMaxRange = 10,\n\t\t\treverbFactor = 0,\n\t\t}}\n\t}}\n}}')
    config.close()

def GetDRUID()->int:
    config = open("./automation/config.txt")
    data = json.load(config)

    BroadCastID = data['BroadcastID']
    config.close()

    with open("./media/lua/client/RadioCom/RadioWavs.lua", "r") as f:
        contents = f.readlines()
    i=1
    for line in contents:
        if BroadCastID in line:
            return i-11        
        if '--END RadioWavs' in line:
            break
        i = i+1
    contents.insert(i-1, 'RadioWavs.files['+ str(i-11) + '] = "' + BroadCastID + '"\n')
    with open("./media/lua/client/RadioCom/RadioWavs.lua", "w") as f:
        contents = "".join(contents)
        f.write(contents)
    return i-11

RadioDataTree = ET.parse('./media/radio/RadioData.xml')
RadioRoot = RadioDataTree.getroot()
config = open("./automation/config.txt")
data = json.load(config)

broadcastID = data['BroadcastID']

config.close()

print(broadcastID)

Broadcast = ""
#Find Broadcast element
for node in RadioRoot.iter('BroadcastEntry'):
    if node.attrib['ID']==broadcastID:
        Broadcast = node

if Broadcast == "":
    print("Broadcast not Found, make sure the ID is the BroadcastEntry ID, not the line Entry, and that it is complete")
    exit()

#Check if first LineEntry element has a DRU code, meaning it is already voiced
if 'codes' in Broadcast[0].attrib:
    if "DRU" in Broadcast[0].attrib['codes']:
        print("Broadcast already voiced. Do you wish to overwrite it? Y/N")
        reply = input()
        reply = reply.upper()
        if 'Y' not in reply:
            print("Quitting")
            exit()

print("Writing Broadcast transcript to tmp file")

N_Voices = WriteBroadcastTranscript(Broadcast)

#Ask user to check the transcript

print("")

print("Please make sure the transcript in automation/tmp/transcript.txt is correct. If you want, you can add any sound effect. Just place the sound effect in /automation/SoundEffects and write YOUR_SFX_NAME in between brackets before or after the voices (and before the script). You can also remove the Voice from the script, and it wont be sent to Eleven Labs. SFX will be concatenated in order of appearance")

AskContinue()

#WriteVoiceConfig(N_Voices,broadcastID)

print("Before continuing make sure the config file is correct, specially the voice names and API key.")
cost = GetCharCost()
print("This will cost " + str(cost) + " Characters")

AskContinue()

SegmentTimes = GetVoiceLines()

#add sound ID to media/script/sounds_RadioWavs.txt
AddSoundIDToScript()
#add sound ID to LUA and get DRU id

dru_id = GetDRUID()

#Write timeTags to RadioData.XML and add DRU code
#Write DRU code
if 'codes' in Broadcast[0].attrib:
    if "DRU" in Broadcast[0].attrib['codes']:
        #make sure dru_id = DRU in XML
        codes = Broadcast[0].attrib['codes'].split(',')
        newCodes = ""
        i=0
        for code in codes:
            if i != 0:
                newCodes += ","
            if 'DRU' in code:
                if int(code[4:]) != dru_id:
                    #change code to match dru_id
                    code = f'DRU+{dru_id}'
            newCodes += code
            i = i+1
        Broadcast[0].set('codes',newCodes)
    else:
        #add DRU code + dru_id
        newCodes = Broadcast[0].attrib['codes'] + f',DRU+{dru_id}'
        Broadcast[0].set('codes',newCodes)
else:
    #add attrib with value DRU+{dru_i}
    Broadcast[0].set('codes',f'DRU+{dru_id}')
#add TimeTags
i=0
for lineEntry in Broadcast:
    #if it has ${t=x.xx} remove it
    Line = lineEntry.text
    if '${' in Line[0:8]:
        Line = Line[9:] #remove ${t:x.xx}
    #add ${t=SegmentTimes}
    lineEntry.text = f'${{t:{SegmentTimes[i]:.2f}}}' + Line
    i = i+1
#write changes to tree

RadioDataTree.write('./media/radio/RadioData.xml')

print(SegmentTimes)

#cleanup, delete files in tmp/lines

for file_name in os.listdir('./automation/tmp/Lines'):
    file = './automation/tmp/Lines/' + file_name
    if os.path.isfile(file):
        os.remove(file)

for file_name in os.listdir('./automation/tmp/'):
    file = './automation/tmp/' + file_name
    if os.path.isfile(file):
        os.remove(file)