from fastapi import FastAPI, Query
from transformers import pipeline

# Create a new FastAPI app instance
app = FastAPI()

# Initialize the text generation pipeline
pipe = pipeline("text2text-generation", model="google/flan-t5-small")



def generate_advice(generated_text):
    # Placeholder function for generating actionable advice based on generated text
    # Replace with actual logic or integrate with medical knowledge base for accurate advice
    # Example keywords for different health issues
    keywords_pain = ["stomach ache", "headache", "pain"]
    keywords_blurred_vision = ["blurred vision", "vision issues", "eye problems"]
    keywords_mental_health = ["anxiety", "depression", "stress"]

    advice = []

    # Check for specific health issues in the generated text
    for keyword in keywords_pain:
        if keyword in generated_text.lower():
            advice.append("Consider taking over-the-counter pain relief medication. If the pain persists or worsens, consult a doctor.")

    for keyword in keywords_blurred_vision:
        if keyword in generated_text.lower():
            advice.append("Schedule an appointment with an ophthalmologist for a thorough eye examination.")

    for keyword in keywords_mental_health:
        if keyword in generated_text.lower():
            advice.append("Practice mindfulness techniques, consider talking to a therapist, or consult a mental health professional for support.")

    # If no specific advice was generated, provide a general recommendation
    if not advice:
        advice.append("Please consult a healthcare professional for a proper diagnosis and treatment.")

    # Add a general prompt suggestion similar to ChatGPT
    advice.append("Feel free to ask more about any specific concerns or questions you have.")

    # Return the advice as a formatted string or list
    return advice



@app.get("/")
def home():
    return {"message": "Hello World"}

# Define a function to handle the GET request at `/generate`
@app.get("/generate")
def generate(text: str = Query(..., title="Input Text", description="Describe your health issue here")):
    # Use the pipeline to generate text based on the input text
    output = pipe(text)
    # Extract the generated text from the pipeline output
    generated_text = output[0]['generated_text']
    # Enhance the response with actionable advice based on the generated text
    advice = generate_advice(generated_text)
    # Return the input text, generated text, and advice as a JSON response
    return {"input": text, "generated_output": generated_text, "advice": advice}