import gradio as gr
import os
from huggingface_hub import InferenceClient
from pypdf import PdfReader
import textwrap

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

# Retrieve the API key from environment variables
api_key = os.getenv("API")

if not api_key:
    raise ValueError("API key not found. Make sure it's set in the environment variables.")

# Instantiate the API client with the API key
client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta", token=api_key)  # Use 'model' instead of 'repo_id'


PROMPT_1 = """Production of Maintenance Schedule

Uploaded here is a vendor Operation and Maintenance manual for a piece of equipment.

First, extract the following information from the document:

*Type of equipment (e.g. pump, valve, lighting fixture, control equipment, high voltage electrical equipment, low voltage electrical equipment).

*Manufacturer's name.

*Vendor's model number.

Secondly, provide a maintenance schedule for this equipment. This schedule should include daily, weekly, monthly, quarterly, annual etc frequency maintenance activities and be presented in a table. Include any specific maintenance activities that are recommended in the manual provided, as well as typical maintenance activities required for this type of equipment.

Note: Order the response by frequency, put non-specific frequencies at the bottom. Do not use the word 'regularly' in the response. Make sure the response includes only information from the document attached.

Here is an example of a maintenance schedule:

Type of equipment: Pump

Manufacturer: Xylem

Model: XB34-U

<table>
<tr>
    <th>ID</th>
    <th>Activity</th>
    <th>Frequency</th>
    <th>Procedure</th>
</tr>
<tr class="even">
    <td>1</td>
    <td>Basic inspection and drive check</td>
    <td>1 Month</td>
    <td>Motors where accessible to be checked monthly for vibration, squeaking, noisy and hot bearings.<br>Belt drives where accessible to be checked for disintegration, fraying or galloping where applicable.<br>Chain drives check lubrication and effective guards where applicable.</td>
</tr>
<tr class="odd">
    <td>2</td>
    <td>Basic pump inspection</td>
    <td>3 Month</td>
    <td>Preparation activity & to obtain a 'Permit to Work' before commencing tasks.<br>Inspect pump for condition. No leaks, unusual noise or vibration.<br>Verify shaft, stator clean and free of mechanical damage.<br>Visually inspect the oil sight gauge to ensure the oil level is adequate.</td>
</tr>
<tr class="even">
    <td>3</td>
    <td>Grease pump bearing</td>
    <td>6 Month</td>
    <td>Check for unusual noise, heat and smell.<br>Grease pump bearing with appropriate lubricant.</td>
</tr>
<tr class="odd">
    <td>4</td>
    <td>Basic inspection and drive check</td>
    <td>1 Month</td>
    <td>Motors where accessible to be checked monthly for vibration, squeaking, noisy and hot bearings.<br>Belt drives where accessible to be checked for disintegration, fraying or galloping where applicable.<br>Chain drives check lubrication and effective guards where applicable.</td>
</tr>
<tr class="even">
    <td>5</td>
    <td>Comprehensive inspection</td>
    <td>12 Month</td>
    <td>Preparation activity & to obtain a 'Permit to Work' before commencing task.<br>Confirm anti condensation heaters are working where installed and heater operating.<br>Lubricate motor bearing with Albida EP2. Refer to motor specifications for lubrication amount.<br>Measure voltage and current in each phase is equal and within rating, each phase within 10% or each other.</td>
</tr>
</table>

<style>
table {
    border-collapse: collapse;
    width: 100%;
}

th, td {
    padding: 8px;
    text-align: left;
    border-bottom: 1px solid #ddd;
}

tr.odd {
    background-color: #f2f2f2;
}

tr.even {
    background-color: #ffffff;
}
</style>
"""

MAX_TOKENS = 1000  # Adjust this value based on your model's maximum token limit


def process_pdf(file_path):
    reader = ''
    reader = PdfReader(file_path)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n"
    
    # Chunk the text into smaller segments
    chunks = textwrap.wrap(text, width=MAX_TOKENS, replace_whitespace=False)
    
    responses = []
    for chunk in chunks:
        response = client.text_generation(f"Document: {chunk}\n" + PROMPT_1)
        if isinstance(response, str):
            generated_text = response
        else:
            response_dict = response.json()
            generated_text = response_dict.get("generated_text", "")
        responses.append(generated_text)
    
    full_response = "\n\n".join(responses)
    table_data = full_response.split("\n\n", 1)[1]
    table_rows = table_data.split("\n")[2:]  # Skip the first two lines (headers)
    table_output = "\n".join(table_rows)  # Convert list to string

    return full_response, table_output

def generate_work_instructions(selected_row, pdf_text):
    if selected_row is None or selected_row == "":
        return "Please select a row from the table."

    row_data = selected_row.split("\t")
    if len(row_data) >= 4:
        activity = row_data[3].strip()
        prompt = f"""You are an engineer who is tasked with developing maintenance work instructions for equipment.

Using only information contained in the attached maintenance manual, write a maintenance work instruction in the style of a technical writer for the {activity} associated with this asset. Format the work instruction as a professional document using Markdown formatting.

The document should have the following structure:

# Maintenance Work Instruction

## Equipment Make and Model
[insert make and model as retrieved from the uploaded document]

## Equipment Type/Classification
[insert equipment type/classification as retrieved from the uploaded document]

## Work Instruction Title
{activity}

### Introduction
• Describe the scope of the work instruction
• Expected work duration

### Safety Precautions
• Describe any hazards to be considered
• How to prepare the work space
• PPE Required
• General instructions on Work Permits
• Pre work inspections
• Isolations

### Tools and Materials
• List of tooling required
• List of materials required. Include technical specifications where it is available from the documents

### Work Steps
• Break the work instruction down into a series of simple steps. Draw upon the text contained within the document
• Provide sub steps as required, with line breaks between each sub step
• Include in the steps how to use the tools and materials as appropriate
• Include the return to service and startup procedures
• Format the work steps as an ordered list with line breaks between each step and sub step

Do not include any further text after these steps above.

Example format:

# Maintenance Work Instruction

## Equipment Make and Model
XYZ Pump

## Equipment Type/Classification
Centrifugal Pump

## Work Instruction Title
Annual Pump Inspection and Maintenance

### Introduction
This work instruction covers the annual inspection and maintenance procedures for the XYZ Centrifugal Pump. The expected work duration is approximately 4 hours.

### Safety Precautions
...

### Tools and Materials
...

### Work Steps
1. Step 1
   - Sub step 1
   - Sub step 2

2. Step 2
   - Sub step 1
   - Sub step 2
   - Sub step 3

...

Document: {pdf_text}"""
        client.user(prompt)
        return client.get_response()
    else:
        return "Invalid row format. Please select a valid row from the table."
    


with gr.Blocks(css=".gradio-container {background-color: #f0f0f0; padding: 20px; border-radius: 10px;}") as demo:
    with gr.Row(elem_id="banner"):
        gr.Markdown("# AmWiz 🧙‍♂️")
    
    with gr.Row():
        with gr.Column():
            pdf_input = gr.File(label="Upload PDF", type="filepath")
            result_output = gr.HTML(label="Maintenance Schedule")
            selected_row = gr.Textbox(label="Select Row (copy and paste from the table)")
            work_instructions_output = gr.Markdown(label="Work Instructions")
            
            pdf_input.upload(process_pdf, inputs=[pdf_input], outputs=[result_output, result_output])
            generate_button = gr.Button("Generate Work Instructions")
            generate_button.click(generate_work_instructions, inputs=[selected_row, pdf_input], outputs=[work_instructions_output])

    with gr.Row(elem_id="footer"):
        gr.Markdown("Made with ❤️ by Asset 2")

demo.launch()