# 导入所需的库
from pycipher import SimpleSubstitution as SimpleSub
import random
import re
from ngram_score import ngram_score
import proability
import gradio as gr

def decrypt_text_internal(ciphertext):
    fitness = ngram_score('quadgrams.txt')
    ctext = re.sub('[^A-Z]', '', ciphertext.upper())
    maxkey = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ')
    maxscore = -99e9
    parentscore, parentkey = maxscore, maxkey[:]
    
    i = 0
    # Limit iterations to prevent infinite loops in a web server environment
    while i < 1000:
        i = i + 1
        random.shuffle(parentkey)
        deciphered = SimpleSub(parentkey).decipher(ctext)
        parentscore = fitness.score(deciphered)
        count = 0
        while count < 1000:
            a = random.randint(0, 25)
            b = random.randint(0, 25)
            child = parentkey[:]
            child[a], child[b] = child[b], child[a]
            deciphered = SimpleSub(child).decipher(ctext)
            score = fitness.score(deciphered)
            if score > parentscore:
                parentscore = score
                parentkey = child[:]
                count = 0
            count = count + 1
        if parentscore > maxscore:
            maxscore, maxkey = parentscore, parentkey[:]
            ss = SimpleSub(maxkey)
            plaintext = ss.decipher(ctext)
            plaintext1 = add_punctuation_and_spaces(ciphertext, plaintext)
            # In a web context, we return the first good result.
            # The original loop was infinite, which is not suitable for a server.
            return plaintext1
    # Fallback if no good solution is found within the iteration limit
    ss = SimpleSub(maxkey)
    plaintext = ss.decipher(ctext)
    return add_punctuation_and_spaces(ciphertext, plaintext)


def output(string1, dic, string2):
    modified_string1 = list(string1)
    modified_string2 = list(string2)
    for i in range(len(string1)):
        if modified_string1[i] in dic and modified_string2[i] != ' ':
            modified_string2[i] = dic[modified_string1[i]]
    modified_string2 = ''.join(modified_string2)
    return modified_string2

def add_punctuation_and_spaces(ciphertext, plaintext):
    result = ""
    j = 0
    for i in range(len(ciphertext)):
        if not ciphertext[i].isalpha():
            result += ciphertext[i]
        else:
            if ciphertext[i].islower():
                result += plaintext[j].lower()
            else:
                # The original code had a bug here, always making it lowercase.
                # This is a guess at the intended behavior.
                result += plaintext[j] 
            j += 1
    return result

def decrypt_interface(ciphertext, key):
    """
    This is the main function that will be exposed through the Gradio interface.
    """
    if not ciphertext:
        return "Please enter some ciphertext."
        
    plaintext = decrypt_text_internal(ciphertext)
    
    if key:
        try:
            key_dic = proability.read_key(key)
            plaintext = output(ciphertext, key_dic, plaintext)
        except Exception as e:
            return f"Error processing key: {e}. Please check the key format (e.g., a=B c=D)."
            
    return plaintext

# Create the Gradio interface
iface = gr.Interface(
    fn=decrypt_interface,
    inputs=[
        gr.Textbox(lines=10, label="Ciphertext", placeholder="Enter the text to decrypt..."),
        gr.Textbox(lines=2, label="Known Key Mappings (Optional)", placeholder="e.g., a=B c=D")
    ],
    outputs=gr.Textbox(lines=10, label="Plaintext"),
    title="Simple Substitution Cipher Decryptor",
    description="An automatic decryption tool for simple substitution ciphers. You can optionally provide known letter mappings to improve accuracy."
)

# Launch the app
if __name__ == "__main__":
    iface.launch(server_name="127.0.0.1", server_port=7860)