import torch
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import HTMLResponse
from transformers import EfficientNetImageProcessor, EfficientNetForImageClassification
from PIL import Image
import io
import sqlite3
import os

app = FastAPI()

# Ensure the db directory exists
DB_DIR = "/app/db"
os.makedirs(DB_DIR, exist_ok=True)

# Load the EfficientNet-B0 model
model_name = "google/efficientnet-b0"
feature_extractor = EfficientNetImageProcessor.from_pretrained(model_name)
model = EfficientNetForImageClassification.from_pretrained(model_name)

# Mapping EfficientNet-B0 ImageNet labels to Bengali fruit names
fruit_mapping = {
    "Granny Smith": "আপেল",
    "banana": "কলা",
    "orange": "কমলা",
    "strawberry": "স্ট্রবেরি",
    "pomegranate": "ডালিম",
    "lemon": "লেবু",
    "pineapple": "আনারস",
    "jackfruit": "কাঠাল"
}

# Health benefits and recommended intake
benefits_mapping = {
    "আপেল": {
        "intake": "১টি মাঝারি আপেল (১৫০ গ্রাম)",
        "diabetic": "১/২টি",
        "benefits": [
            "হজম সমস্যা: কোষ্ঠকাঠিন্য ও IBS উপশম করে",
            "ডায়াবেটিস নিয়ন্ত্রণ: রক্তে শর্করার মাত্রা স্থিতিশীল রাখে",
            "হার্ট ডিজিজ: LDL কোলেস্টেরল ১০-১৫% কমায়",
            "ওজন নিয়ন্ত্রণ: মেটাবলিক সিন্ড্রোম রোধ করে"
        ]
    },
    "কলা": {
        "intake": "১-২টি মাঝারি কলা (১০০-১৫০ গ্রাম)",
        "diabetic": "১টি",
        "benefits": [
            "উচ্চ রক্তচাপ: দিনে ২টি কলা স্ট্রোকের ঝুঁকি ২১% কমায়",
            "ডিপ্রেশন: সেরোটোনিন উৎপাদন বাড়িয়ে মুড উন্নত করে",
            "হজম সমস্যা: আলসার ও গ্যাস্ট্রাইটিসে উপকারী",
            "ক্রীড়াবিদদের জন্য: মাসল ক্র্যাম্প প্রতিরোধ করে"
        ]
    },
    "কমলা": {
        "intake": "১টি বড় কমলা বা ১ গ্লাস জুস",
        "diabetic": "১/২টি",
        "benefits": [
            "ইমিউনিটি: সাধারণ সর্দি-কাশির স্থায়িত্ব ২৩% কমায়",
            "কিডনি স্টোন: সাইট্রেট স্টোন গলাতে সাহায্য করে",
            "রক্তাল্পতা: আয়রন শোষণ ৬৭% বাড়ায়",
            "ত্বকের ক্যান্সার: UV রশ্মির ক্ষতি থেকে রক্ষা করে"
        ]
    },
    "স্ট্রবেরি": {
        "intake": "৫-৬টি স্ট্রবেরি (১০০ গ্রাম)",
        "diabetic": "১-২টি",
        "benefits": [
            "ক্যান্সার: টিউমার বৃদ্ধি ৫০% ধীর করে",
            "গর্ভাবস্থা: নিউরাল টিউব ডিফেক্ট প্রতিরোধ করে",
            "আর্থ্রাইটিস: জয়েন্টের ব্যথা ৩০% কমায়",
            "চোখের ছানি: অক্সিডেটিভ স্ট্রেস কমায়"
        ]
    },
    "ডালিম": {
        "intake": "১টি মাঝারি ডালিম (১৫০ গ্রাম) বা ১/২ কাপ দানা",
        "diabetic": "১/৪ কাপ দানা",
        "benefits": [
            "হৃদরোগ প্রতিরোধ: রক্তচাপ কমায়, ধমনীর স্থিতিস্থাপকতা বাড়ায়",
            "রক্তস্বল্পতা দূর করে: আয়রন ও ভিটামিন সি হিমোগ্লোবিন বৃদ্ধি করে",
            "ক্যান্সার প্রতিরোধ: প্রোস্টেট ও ব্রেস্ট ক্যান্সার",
            "গাঁটের ব্যথা কমায়: Arthritis উপশম করে",
            "মস্তিষ্কের স্বাস্থ্য: আলঝেইমার্স রোগের ঝুঁকি কমায়"
        ]
    },
    "লেবু": {
        "intake": "১টি মাঝারি লেবু বা ১/২ গ্লাস জুস",
        "diabetic": "১/৪টি",
        "benefits": [
            "ইমিউনিটি: ভিটামিন সি দিয়ে রোগ প্রতিরোধ ক্ষমতা বাড়ায়",
            "হজম: অ্যাসিডিটি ও বদহজম কমায়",
            "ত্বকের যত্ন: অ্যান্টিঅক্সিডেন্ট দিয়ে ত্বক উজ্জ্বল করে",
            "ওজন কমানো: মেটাবলিজম বাড়াতে সাহায্য করে"
        ]
    },
    "আনারস": {
        "intake": "১ কাপ কাটা আনারস (১৫০ গ্রাম)",
        "diabetic": "১/২ কাপ",
        "benefits": [
            "হজম: ব্রোমেলিন এনজাইম দিয়ে প্রোটিন ভাঙতে সাহায্য করে",
            "প্রদাহ কমায়: আর্থ্রাইটিস ও পেশির ব্যথা উপশম করে",
            "ইমিউনিটি: ভিটামিন সি দিয়ে সংক্রমণ প্রতিরোধ করে",
            "দাঁতের স্বাস্থ্য: মাড়ির প্রদাহ কমায়"
        ]
    },
    "কাঠাল": {
        "intake": "১ কাপ কাটা কাঠাল (১৫০ গ্রাম)",
        "diabetic": "কম পরিমাণে",
        "benefits": [
            "শক্তি: উচ্চ কার্বোহাইড্রেট দিয়ে শক্তি বাড়ায়",
            "ইমিউনিটি: ভিটামিন এ ও সি দিয়ে রোগ প্রতিরোধ করে",
            "হৃদরোগ: পটাশিয়াম দিয়ে রক্তচাপ নিয়ন্ত্রণ করে",
            "কোষ্ঠকাঠিন্য: ফাইবার দিয়ে হজমশক্তি বাড়ায়"
        ]
    }
}

# SQLite database initialization for nutrition
def init_db():
    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_nutrition.db"))
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS fruits (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE,
                    vitamin TEXT,
                    mineral TEXT,
                    carbohydrate TEXT,
                    protein TEXT,
                    amino_acid TEXT)''')
    fruits_data = [
        ("আপেল", "ভিটামিন সি, ভিটামিন এ", "পটাশিয়াম, ম্যাগনেসিয়াম", "১৪ গ্রাম", "০.৩ গ্রাম", "নেই"),
        ("কলা", "ভিটামিন বি৬, ভিটামিন সি", "পটাশিয়াম", "২৭ গ্রাম", "১.১ গ্রাম", "ট্রিপটোফ্যান"),
        ("কমলা", "ভিটামিন সি", "ক্যালসিয়াম, পটাশিয়াম", "১২ গ্রাম", "০.৯ গ্রাম", "নেই"),
        ("স্ট্রবেরি", "ভিটামিন সি, ভিটামিন কে", "ম্যাঙ্গানিজ, ফোলেট", "৮ গ্রাম", "০.৭ গ্রাম", "নেই"),
        ("ডালিম", "ভিটামিন সি", "আয়রন, ক্যালসিয়াম", "১৯ গ্রাম", "১.৭ গ্রাম", "নেই"),
        ("লেবু", "ভিটামিন সি", "পটাশিয়াম, ম্যাগনেসিয়াম", "৯ গ্রাম", "১.১ গ্রাম", "নেই"),
        ("আনারস", "ভিটামিন সি, ভিটামিন বি৬", "ম্যাঙ্গানিজ, পটাশিয়াম", "২২ গ্রাম", "০.৫ গ্রাম", "নেই"),
        ("কাঠাল", "ভিটামিন এ, ভিটামিন সি", "পটাশিয়াম, ম্যাগনেসিয়াম", "৩৮ গ্রাম", "১.৭ গ্রাম", "নেই")
    ]
    c.executemany("INSERT OR IGNORE INTO fruits (name, vitamin, mineral, carbohydrate, protein, amino_acid) VALUES (?, ?, ?, ?, ?, ?)", fruits_data)
    conn.commit()
    conn.close()

# New database for diabetic-specific information
def init_diabetes_db():
    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_diabetes.db"))
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS diabetes_info (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    fruit_name TEXT UNIQUE,
                    glycemic_index TEXT,
                    portion_size TEXT,
                    best_time TEXT,
                    precautions TEXT)''')
    
    diabetes_data = [
        ("আপেল", "নিম্ন (৩৮)", "১/২টি মাঝারি", "সকাল বা দুপুর", "অতিরিক্ত খাবেন না"),
        ("কলা", "মাঝারি (৫১)", "১টি ছোট", "সকালে", "পাকা কলা এড়িয়ে চলুন"),
        ("কমলা", "নিম্ন (৪৩)", "১/২টি", "দুপুরে", "জুস এড়িয়ে ফল খান"),
        ("স্ট্রবেরি", "নিম্ন (৪০)", "১-২টি", "যেকোনো সময়", "পরিমিত খান"),
        ("ডালিম", "নিম্ন (৫৩)", "১/৪ কাপ দানা", "সকালে", "বীজ সহ খান"),
        ("লেবু", "খুব নিম্ন (২০)", "১/৪টি", "যেকোনো সময়", "চিনি ছাড়া ব্যবহার করুন"),
        ("আনারস", "মাঝারি (৫৯)", "১/২ কাপ", "সকালে", "পরিমিত খান"),
        ("কাঠাল", "মাঝারি (৫৭)", "২-৩টি কোয়া", "দুপুরে", "অতিরিক্ত এড়িয়ে চলুন")
    ]
    
    c.executemany("INSERT OR IGNORE INTO diabetes_info (fruit_name, glycemic_index, portion_size, best_time, precautions) VALUES (?, ?, ?, ?, ?)", diabetes_data)
    conn.commit()
    conn.close()

init_db()
init_diabetes_db()

# Serve the HTML file
@app.get("/", response_class=HTMLResponse)
async def home():
    with open("index.html", "r", encoding="utf-8") as f:
        return f.read()

# Prediction endpoint
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    predicted_class_idx = outputs.logits.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_class_idx].split(",")[0].strip()
    bengali_fruit = fruit_mapping.get(predicted_label, predicted_label)

    # Fetch nutrition from database
    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_nutrition.db"))
    c = conn.cursor()
    c.execute("SELECT vitamin, mineral, carbohydrate, protein, amino_acid FROM fruits WHERE name = ?", (bengali_fruit,))
    nutrition = c.fetchone()
    conn.close()
    if not nutrition:
        nutrition = ("অজানা", "অজানা", "অজানা", "অজানা", "অজানা")

    # Fetch health benefits
    benefits_info = benefits_mapping.get(bengali_fruit, {
        "intake": "অজানা",
        "diabetic": "অজানা",
        "benefits": ["তথ্য নেই"]
    })

    # Fetch diabetes info
    conn = sqlite3.connect(os.path.join(DB_DIR, "fruit_diabetes.db"))
    c = conn.cursor()
    c.execute("SELECT glycemic_index, portion_size, best_time, precautions FROM diabetes_info WHERE fruit_name = ?", (bengali_fruit,))
    diabetes_info = c.fetchone()
    conn.close()
    if not diabetes_info:
        diabetes_info = ("অজানা", "অজানা", "অজানা", "অজানা")

    return {
        "filename": file.filename,
        "predicted_class": bengali_fruit,
        "nutrition": {
            "vitamin": nutrition[0],
            "mineral": nutrition[1],
            "carbohydrate": nutrition[2],
            "protein": nutrition[3],
            "amino_acid": nutrition[4]
        },
        "health": {
            "intake": benefits_info["intake"],
            "diabetic": benefits_info["diabetic"],
            "benefits": benefits_info["benefits"]
        },
        "diabetes": {
            "glycemic_index": diabetes_info[0],
            "portion_size": diabetes_info[1],
            "best_time": diabetes_info[2],
            "precautions": diabetes_info[3]
        }
    }