import sqlite3
import gradio as gr
from hashlib import md5 as hash_algo
from re import match
from io import BytesIO
from pypdf import PdfReader
from llm_rs import AutoModel,SessionConfig,GenerationConfig,Precision

repo_name = "rustformers/mpt-7b-ggml"
file_name = "mpt-7b-instruct-q5_1-ggjt.bin"
script_env = 'prod'


session_config = SessionConfig(threads=2,batch_size=2)
model = AutoModel.from_pretrained(repo_name, model_file=file_name, session_config=session_config,verbose=True)

def process_stream(rules, log, temperature, top_p, top_k, max_new_tokens, seed):
    con = sqlite3.connect("history.db")
    cur = con.cursor()
    instruction = ''
    hashes = []
    
    if type(rules) is not list:
        rules = [rules]
    
    for rule in rules:
        data, hash = get_file_contents(rule)
        instruction += data + '\n'
        hashes.append(hash)
    
    hashes.sort()
    hashes = hash_algo(''.join(hashes).encode()).hexdigest()
    
    largest = 0
    lines = instruction.split('\r\n')
    
    if len(lines) == 1:
        lines = instruction.split('\n')
    
    for line in lines:
        m = match('^(\d+)\.', line)
        if m != None:
            num = int(line[m.start():m.end()-1])
            
            if num > largest:
                largest = num
                
    instruction += str(largest + 1) + '. '
    
    query, hash = get_file_contents(log)
    hashes = hash_algo((hashes + hash).encode()).hexdigest()
    
    instruction = instruction.replace('\r\r\n', '\n')

    full_req = "A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.\r\n\r\nQ: Read the rules stated below and check the queries for any violation. State the rules which are violated by a query (if any). Also suggest a possible remediation, if possible. Do not make any assumptions outside of the rules stated below.\r\n\r\n" + instruction + 'The queries are as follows:\r\n' + query + '\r\n  \r\nA: '
    full_req = full_req.replace('\r\n', '\n')
    
    prompt=f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{full_req}
### Response:
Answer:"""
    
    response = ""
    row = cur.execute('SELECT response FROM queries WHERE hexdigest = ?', [hashes]).fetchone()

    if row != None:
        response += "Cached Result:\n" + row[0]
        yield response
    else:
        if script_env != 'test':
            generation_config = GenerationConfig(seed=seed,temperature=temperature,top_p=top_p,top_k=top_k,max_new_tokens=max_new_tokens)
            streamer = model.stream(prompt=prompt,generation_config=generation_config)
            for new_text in streamer:
                response += new_text
                yield response
        else:
            num = 0
            while num < 100:
                response += " " + str(num)
                num += 1
                yield response
                
        cur.execute('INSERT INTO queries VALUES(?, ?)', (hashes, response))
        con.commit()
        
    cur.close()
    con.close()

def get_file_contents(file):
    data = None
    byte_hash = ''
    
    with open(file.name, 'rb') as f:
        data = f.read()
        byte_hash = hash_algo(data).hexdigest()
    
    if file.name.endswith('.pdf'):
        rdr = PdfReader(BytesIO(data))
        data = ''
        
        for page in rdr.pages:
            data += page.extract_text()
    else:            
        data = data.decode()
        
        if file.name.endswith(".csv"):
            data = data.replace(',', ' ')
            
    return (data, byte_hash)

def upload_log_file(files):
    file_paths = [file.name for file in files]
    return file_paths
    
def upload_file(files):
    file_paths = [file.name for file in files]
    return file_paths

with gr.Blocks(
    theme=gr.themes.Soft(),
    css=".disclaimer {font-variant-caps: all-small-caps;}",
) as demo:
    gr.Markdown(
        """<h1><center>Grid 5.0 Information Security Track</center></h1>
        """
    )
    
    rules = gr.File(file_count="multiple")
    upload_button = gr.UploadButton("Click to upload a new Compliance Document", file_types=[".txt", ".pdf"], file_count="multiple")
    upload_button.upload(upload_file, upload_button, rules)
    
    with gr.Row():
        with gr.Column():
            log = gr.File()
            upload_log_button = gr.UploadButton("Click to upload a log file", file_types=[".txt", ".csv", ".pdf"], file_count="multiple")
            upload_log_button.upload(upload_log_file, upload_log_button, log)

            with gr.Accordion("Advanced Options:", open=False):
                with gr.Row():
                    with gr.Column():
                        with gr.Row():
                            temperature = gr.Slider(
                                label="Temperature",
                                value=0.8,
                                minimum=0.1,
                                maximum=1.0,
                                step=0.1,
                                interactive=True,
                                info="Higher values produce more diverse outputs",
                            )
                    with gr.Column():
                        with gr.Row():
                            top_p = gr.Slider(
                                label="Top-p (nucleus sampling)",
                                value=0.95,
                                minimum=0.0,
                                maximum=1.0,
                                step=0.01,
                                interactive=True,
                                info=(
                                    "Sample from the smallest possible set of tokens whose cumulative probability "
                                    "exceeds top_p. Set to 1 to disable and sample from all tokens."
                                ),
                            )
                    with gr.Column():
                        with gr.Row():
                            top_k = gr.Slider(
                                label="Top-k",
                                value=40,
                                minimum=5,
                                maximum=80,
                                step=1,
                                interactive=True,
                                info="Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.",
                            )
                    with gr.Column():
                        with gr.Row():
                            max_new_tokens = gr.Slider(
                                label="Maximum new tokens",
                                value=256,
                                minimum=0,
                                maximum=1024,
                                step=5,
                                interactive=True,
                                info="The maximum number of new tokens to generate",
                            )

                    with gr.Column():
                        with gr.Row():
                            seed = gr.Number(
                                label="Seed",
                                value=42,
                                interactive=True,
                                info="The seed to use for the generation",
                                precision=0
                            )
    with gr.Row():
        submit = gr.Button("Submit")
    with gr.Row():
        with gr.Box():
            gr.Markdown("**Output**")
            output_7b = gr.Markdown()

    submit.click(
        process_stream,
        inputs=[rules, log, temperature, top_p, top_k, max_new_tokens,seed],
        outputs=output_7b,
    )

demo.queue(max_size=4, concurrency_count=1).launch(debug=True)