import requests
import json
import os
import transformers
from transformers import ViTFeatureExtractor, ViTForImageClassification
import warnings
from PIL import Image
from io import BytesIO
warnings.filterwarnings('ignore')
api_key = os.environ.get('API_KEY')
import gradio as gr

model_name = "google/vit-base-patch16-224"
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)

def food_identification(img_path):
    try:
        image = Image.open(img_path)
        input = feature_extractor(images= image,return_tensors= 'pt' ) 
        outputs = model(**input)
        logits = outputs.logits
        # model predicts one of the 1000 ImageNet classes
        pred_class_idx = logits.argmax(-1).item()
        ans = f"{model.config.id2label[pred_class_idx]}".split(',')[0]
        return ans
    except FileNotFoundError as e:
        print('Got an error -> ', e)
        return None
    

def cal_calculation(fruit_name):
    api_url = f'https://api.api-ninjas.com/v1/nutrition?query={fruit_name}'
    Headers = { 'X-Api-Key': api_key}
    response = requests.get(api_url, headers=Headers)
    # we are sending api_key through header because header is more secure as compare to body or api_url .
    if response.status_code == 200:
        return response.json()
    else:
        return f"Error: {response.status_code}, Messages: {response.text}"


def json_to_html_table(json_data):
    # Start the HTML table
    html = '<table border="1" style="border-collapse: collapse; width: 50%;">'
    
    # Add table rows
    for item in json_data:
        for key, value in item.items():
            html += '<tr>'
            html += f'<th style="padding: 10px; background-color: #f2f2f2; color: blue; text-align: left;">{key}</th>'
            html += f'<td style="padding: 10px; text-align: left;">{value}</td>'
            html += '</tr>'
    
    # End the HTML table
    html += '</table>'
    return html

def display_table(json_input):
    json_data = json.loads(json_input)
    return json_to_html_table(json_data)


def main_function(image_path):
    try:
        food_identity = food_identification(image_path)
        if not food_identity:
            return "<p style='color:red;'>Could not identify the food item.</p>"
        nutrition_info = cal_calculation(food_identity)
        if isinstance(nutrition_info, str):  # Error message from API
            return f"<p style='color:red;'>{nutrition_info}</p>"
        json_data = json.dumps(nutrition_info)
        return display_table(json_data)
        
    except Exception as e:
        return f"<p style='color:red;'>Unexpected Error: {str(e)}</p>"

import gradio as gr
iface = gr.Interface(
    fn= main_function,
    inputs= gr.Image(type='filepath'),
    outputs= 'html',
    title= 'Food Identification and neutrition info',
    description= 'Upload the image of a food to get neutrition info',
    allow_flagging='never',
    
)

iface.launch(share= True)