import os
import torch
from transformers import pipeline
import gradio as gr
import asyncio
import ipaddress
from typing import Tuple
from accelerate import Accelerator

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

accelerator = Accelerator()

gpt2_pipeline = accelerator.prepare(
    pipeline("text-generation", model="Qwen/Qwen-1_8B-Chat", device=accelerator.device , trust_remote_code=True)
)
Najeb_pipeline = accelerator.prepare(
    pipeline("text-generation", model="najeebjust/Najeeb", device=accelerator.device , trust_remote_code=True)
)
llama2_pipeline = accelerator.prepare(
    pipeline("text-generation", model="Harikrishnan46624/finetuned_llama2-1.1b-chat", device=accelerator.device , trust_remote_code=True)
)
'''
gpt2_pipeline = pipeline("text-generation", model="Qwen/Qwen-1_8B-Chat", device=0 if torch.cuda.is_available() else -1, trust_remote_code=True)
Najeb_pipeline = pipeline("text-generation", model="najeebjust/Najeeb", device=0 if torch.cuda.is_available() else -1)
llama2_pipeline = pipeline("text-generation", model="Harikrishnan46624/finetuned_llama2-1.1b-chat", device=0 if torch.cuda.is_available() else -1)
'''
summarization_pipeline = pipeline("summarization", model="Falconsai/text_summarization", device=0 if torch.cuda.is_available() else -1)


previous_questions = []


async def generate_gpt2(question, max_length, num_beams, temperature):
    return gpt2_pipeline(
        question,
        max_length=max_length,
        num_return_sequences=1,
        num_beams=num_beams,
        do_sample=True,
        top_k=30,
        top_p=0.9,
        temperature=temperature
    )[0]['generated_text']

async def generate_Najeb(question, max_length, num_beams, temperature):
    return Najeb_pipeline(
        question,
        max_length=max_length,
        num_return_sequences=1,
        num_beams=num_beams,
        do_sample=True,
        top_k=30,
        top_p=0.85,
        temperature=temperature
    )[0]['generated_text']

async def generate_llama2(question, max_length, num_beams, temperature):
    return llama2_pipeline(
        question,
        max_length=max_length,
        num_return_sequences=1,
        num_beams=num_beams,
        do_sample=True,
        top_k=30,
        top_p=0.9,
        temperature=temperature
    )[0]['generated_text']


async def generate_responses_async(question, max_length=128, num_beams=2, temperature=0.5):
    responses = {}


    previous_questions.append(question)

    gpt2_task = asyncio.create_task(generate_gpt2(question, max_length, num_beams, temperature))
    Najeb_task = asyncio.create_task(generate_Najeb(question, max_length, num_beams, temperature))
    llama2_task = asyncio.create_task(generate_llama2(question, max_length, num_beams, temperature))

    gpt2_response, Najeb_response, llama2_response = await asyncio.gather(gpt2_task, Najeb_task, llama2_task)

    responses['GPT-2'] = gpt2_response
    responses['Najeb '] = Najeb_response
    responses['LLaMA 2'] = llama2_response


    combined_responses = f"GPT-2: {gpt2_response}\nNajeb: {Najeb_response}\nLLaMA 2: {llama2_response}"


    summarized_response = summarization_pipeline(combined_responses, max_length=150, min_length=50, do_sample=False)[0]['summary_text']

    return {

        "Najeb Answering Response": Najeb_response,
        "GPT-2 Answering Response": gpt2_response,
        "LLaMA 2 Answering Response": llama2_response,
        "Summarized Answering Response": summarized_response,
        "Previous Questions": "\n".join(previous_questions[-5:])
    }


def get_network(ip_input: str) -> Tuple[ipaddress.IPv4Network, str]:
    try:
        if ip_input.count("/") == 0:
            ip_input += "/24"
        net = ipaddress.IPv4Network(ip_input, strict=False)
        ip = ip_input.split("/")[0]
        return (net, ip)
    except ValueError:
        return None, None

def calculate_subnet(ip_input: str) -> str:
    network, ip = get_network(ip_input)
    if network is None or ip is None:
        return "Invalid IP Address or Subnet!"


    network_address = network.network_address
    broadcast_address = network.broadcast_address
    usable_hosts = list(network.hosts())
    num_usable_hosts = len(usable_hosts)
    usable_hosts_range = f"{usable_hosts[0]} - {usable_hosts[-1]}" if usable_hosts else "NA"


    octets = str(ip).split('.')
    binary_octets = [bin(int(octet))[2:].zfill(8) for octet in octets]
    bin_ip = '.'.join(binary_octets)

    bin_addr = str(bin(int(network_address))[2:].zfill(32))
    bin_addr = '.'.join([bin_addr[i:i+8] for i in range(0, len(bin_addr), 8)])

    bin_mask = str(bin(int(network.netmask))[2:].zfill(32))
    bin_mask = '.'.join([bin_mask[i:i+8] for i in range(0, len(bin_mask), 8)])


    result = f"""
IP Address:             {ip}
Address (bin):          {bin_ip}
Network Address:        {network_address}
Network Address (bin):  {bin_addr}
Netmask:                {network.netmask}
Netmask (bin):          {bin_mask}
CIDR Notation:          {network.prefixlen}
Broadcast Address:      {broadcast_address}
Usable IP Range:        {usable_hosts_range}
Number of Hosts:        {network.num_addresses:,d}
Number of Usable Hosts: {num_usable_hosts:,d}
Wildcard Mask:          {network.hostmask}
Private IP:             {network.is_private}
"""
    return result.strip()



def handle_mode_selection(mode, input_text, max_length, num_beams, temperature):
    if mode == "AI Question Answering":
        result = asyncio.run(generate_responses_async(input_text, max_length, num_beams, temperature))
        return result, ""
    else:
        subnet_result = calculate_subnet(input_text)
        return {"Subnet Calculation Result": subnet_result}, ""


custom_css = """
body {
    background-color: #f0f8ff;
    font-family: 'Arial', sans-serif;
    color: #333;
}

h1 {
    text-align: center;
    color: #0066cc;
}

p {
    text-align: center;
    color: #333;
}

.gradio-container {
    width: 80%;
    margin: auto;
    background-color: rgba(255, 255, 255, 0.8);
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
    padding: 20px;
    border-radius: 10px;
}

.gr-button {
    background-color: #0066cc;
    color: white;
    border: none;
    border-radius: 5px;
    padding: 10px;
    cursor: pointer;
    transition: background-color 0.3s ease;
}

.gr-button:hover {
    background-color: #004c99;
}

.gr-textbox {
    border: 2px solid #0066cc;
    border-radius: 5px;
    padding: 10px;
    background-color: #fff;
    color: #333;
}

.gr-slider {
    color: #0066cc;
}

.gr-json {
    background-color: rgba(240, 248, 255, 0.8);
    border-radius: 10px;
    padding: 10px;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
}


#image-container {
    text-align: center;
    position: relative;
}

#image-container img {
    width: 1400px;
    border-radius: 10px;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}

#image-container button {
    position: absolute;
    top: 50%;
    left: 50%;
    transform: translate(-50%, -50%);
    background-color: rgba(0, 102, 204, 0.8);
    color: white;
    border: none;
    padding: 10px 20px;
    border-radius: 5px;
    cursor: pointer;
    font-size: 16px;
    transition: background-color 0.3s ease;
}

#image-container button:hover {
    background-color: rgba(0, 76, 153, 0.8);
}
"""


scroll_js = """
<script>
    function scrollToTop() {
        document.getElementById('target-section').scrollIntoView({behavior: 'smooth'});
    }
</script>
"""


iface = gr.Blocks(css=custom_css)

with iface:
    gr.Markdown(f"<h1>Welcome to Najeb</h1><p>AI Question & Subnet Calculator, Enter your question or IP address to generate answers or calculate subnets.</p>")


    gr.HTML(f"""
    <div id="image-container">
        <img src="https://news.cornell.edu/sites/default/files/styles/story_thumbnail_xlarge/public/2024-07/robot-1280x720_0.jpg?itok=AF6MakCq" alt="AI Image">
        <button onclick="scrollToTop()">Go to Top</button>
    </div>
    {scroll_js}  <!-- Adding the JS to handle scrolling -->
    """)


    gr.Markdown("<div id='target-section'></div>")

    with gr.Row():
        mode_selector = gr.Radio(["AI Question Answering", "Subnet Calculation"], label="Select Mode", value="AI Question Answering")

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Enter your question or IP", placeholder="Type here...", lines=2)
            max_length_slider = gr.Slider(minimum=50, maximum=1024, value=128, label="Max Length")
            num_beams_slider = gr.Slider(minimum=1, maximum=10, value=2, label="Number of Beams", step=1)
            temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.5, label="Temperature", step=0.1)
            submit_button = gr.Button("Submit")

        with gr.Column():
            output_box = gr.JSON(label="Response Output")
            previous_questions_box = gr.Markdown("### Previous Questions\n")

    submit_button.click(
        handle_mode_selection,
        inputs=[mode_selector, input_text, max_length_slider, num_beams_slider, temperature_slider],
        outputs=[output_box, previous_questions_box]
    )


iface.launch(share=True)