import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the Hugging Face Inference Client
client = InferenceClient("Qwen/Qwen2.5-Coder-32B-Instruct")

# Function to stream the compliance suggestions as they are generated
def analyze_compliance_stream(code, compliance_standard, consent):
    if not consent:
        return "You must consent to the processing of your code snippet."

    prompt = f"Analyze the following code for {compliance_standard} compliance and suggest modifications or refactoring to meet the guidelines:\n\n{code}"
    
    messages = [
        {"role": "user", "content": prompt}
    ]

    # Create a stream to receive generated content
    stream = client.chat.completions.create(
        messages=messages,
        temperature=0.5,
        max_tokens=32000,
        top_p=0.7,
        stream=True
    )

    # Stream content as it is generated
    compliance_suggestions = ""
    for chunk in stream:
        compliance_suggestions += chunk.choices[0].delta.content
        yield compliance_suggestions  # Yield incremental content to display immediately

    # Delete the code snippet after processing
    del code

# Create Gradio interface with the modified layout
with gr.Blocks() as app:
    gr.Markdown("## Code Compliance Advisor")
    gr.Markdown("Analyze your code for legal compliance and security standards (e.g., GDPR, HIPAA) and receive actionable suggestions.")
    gr.Markdown("### Privacy Notice")
    gr.Markdown("""
    By using this tool, you consent to the processing of your code snippet for compliance analysis. Your code will be sent to a third-party model for analysis and the results will be displayed to you. We do not store your code or the results unless required by law. 
    You have the right to access and erase your data. If you wish to request access to your data or have it erased, please contact us at [contact@example.com].
    """)

    with gr.Row():
        # First column for input components
        with gr.Column():
            code_input = gr.Textbox(lines=10, label="Code Snippet", placeholder="Enter your code here", elem_id="full_width")
            compliance_standard = gr.Dropdown(
                choices=["GDPR", "HIPAA", "PCI-DSS", "SOC 2", "ISO 27001"],
                label="Compliance Standard",
                value="GDPR"
            )
            consent_checkbox = gr.Checkbox(label="I consent to the processing of my code snippet", value=False)
            analyze_button = gr.Button("Analyze Compliance")

        # Second column for output
        with gr.Column():
            gr.Markdown("### Compliance Suggestions")  # This acts as the label for the output
            output_markdown = gr.Markdown()

    # Link button to function with inputs and outputs
    analyze_button.click(fn=analyze_compliance_stream, inputs=[code_input, compliance_standard, consent_checkbox], outputs=output_markdown)

# Run the Gradio app
app.launch()