from transformers import BlipProcessor, BlipForConditionalGeneration
from sentence_transformers import SentenceTransformer
import gradio as gr
from PIL import Image
import torch
import requests
from bs4 import BeautifulSoup

# Load BLIP Model
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Map common fish names to FishBase scientific names
name_map = {
    "pufferfish": "Tetraodon",
    "stonefish": "Synanceia",
    "lionfish": "Pterois",
    "tuna": "Thunnus",
    "salmon": "Salmo-salar",
    "catfish": "Ictalurus",
    "tilapia": "Oreochromis"
}

# Poisonous species (scientific names)
poisonous_species = ["Tetraodon", "Synanceia", "Pterois"]

# FishBase scraping function
def get_fishbase_summary(scientific_name):
    search_url = f"https://www.fishbase.se/summary/{scientific_name}.html"
    try:
        response = requests.get(search_url, timeout=10)
        if response.status_code != 200:
            return f"FishBase entry not found for: {scientific_name}"

        soup = BeautifulSoup(response.text, "html.parser")
        summary_section = soup.find("div", {"id": "ssummary"})
        if summary_section:
            paragraphs = summary_section.find_all("p")
            text = "\n\n".join(p.get_text(strip=True) for p in paragraphs if p.get_text(strip=True))
            return text or f"No summary available for {scientific_name}"
        else:
            return f"No detailed summary found for {scientific_name}"

    except Exception as e:
        return f"Error fetching FishBase data for {scientific_name}: {str(e)}"

# Fish identification function
def identify_fish(image):
    # Step 1: Generate caption from image
    inputs = blip_processor(image, return_tensors="pt")
    out = blip_model.generate(**inputs)
    caption = blip_processor.decode(out[0], skip_special_tokens=True)

    # Step 2: Extract fish name from caption
    fish_name = None
    for name in name_map:
        if name in caption.lower():
            fish_name = name
            break

    if not fish_name:
        return f"❌ Could not identify a known fish species in the image caption: '{caption}'"

    # Step 3: Lookup in FishBase
    scientific_name = name_map[fish_name]
    summary = get_fishbase_summary(scientific_name)

    # Step 4: Check toxicity
    is_poisonous = "Yes 🧪" if scientific_name in poisonous_species else "No ✅"

    # Step 5: Final Output
    return f"**Image Caption:** {caption}\n\n**Detected Fish:** {fish_name.title()}\n**Scientific Name:** {scientific_name}\n**Poisonous:** {is_poisonous}\n\n**📚 FishBase Info:**\n{summary}"

# Gradio UI
demo = gr.Interface(
    fn=identify_fish,
    inputs=gr.Image(type="pil"),
    outputs="markdown",
    title="🐟 Smart Fish Identifier (BLIP + FishBase)",
    description="Upload a fish image. We use BLIP to describe the fish, match it with known species, then fetch info from FishBase to check if it's poisonous."
)

if __name__ == '__main__':
    demo.launch()