from transformers import pipeline
from PIL import Image
import base64
import io
import json
import argparse
import os
import warnings
import sys

# Suppress warnings
warnings.filterwarnings("ignore")

# === Load model ===
print("🔄 Loading model...")

try:
    model = pipeline(
        "image-classification",
        model="linkanjarad/mobilenet_v2_1.0_224-plant-disease-identification",
        top_k=1
    )
except Exception as e:
    print(f"❌ Failed to load model: {e}")
    sys.exit(1)

print("✅ Model loaded.")

# === Mapping file loader ===
def load_map(file_path):
    mapping = {}
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                if ":" in line:
                    k, v = line.strip().split(":", 1)
                    mapping[k.strip()] = v.strip()
    except Exception as e:
        print(f"⚠️ Failed to load {file_path}: {e}")
    return mapping

# === Load mappings ===
disease_map = load_map("diseases.txt")
treatment_map = load_map("treatments.txt")
fertilizer_map = load_map("fertilizers.txt")

# === Load critical diseases ===
try:
    with open("critical_diseases.txt", "r", encoding="utf-8") as f:
        critical_diseases = set(line.strip() for line in f if line.strip())
except Exception as e:
    print(f"⚠️ Failed to load critical_diseases.txt: {e}")
    critical_diseases = set()

# === Prediction Function ===
def predict_disease(base64_img):
    try:
        img_bytes = base64.b64decode(base64_img)
        image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    except Exception as e:
        return {"error": f"Invalid image input: {str(e)}"}

    try:
        result = model(image)[0]
        label = result["label"]
        confidence = float(result["score"])
        disease = disease_map.get(label, label)
        treatment = treatment_map.get(disease, "Consult an expert.")
        fertilizer = fertilizer_map.get(disease, "Use general NPK 10-10-10.")

        output = {
            "disease_prediction": disease,
            "confidence": round(confidence, 4),
            "suggested_treatment": treatment,
            "fertilizer_recommendation": fertilizer
        }

        if confidence < 0.6 or disease in critical_diseases:
            output["alert"] = True

        return output
    except Exception as e:
        return {"error": f"Prediction failed: {str(e)}"}

# === Main CLI ===
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="🌿 Plant Disease Detector\n\nUsage examples:\n  python app.py --image leaf.jpg --raw\n  python app.py --image encoded.txt",
        formatter_class=argparse.RawTextHelpFormatter
    )

    parser.add_argument("--image", type=str, help="Path to image or base64 file", required=False)
    parser.add_argument("--raw", action="store_true", help="Use this flag if the input is a raw image file (e.g., JPG, PNG)")
    args = parser.parse_args()

    if not args.image:
        print("❗ Error: You must provide an image file using --image\n")
        parser.print_help()
        sys.exit(1)

    if not os.path.isfile(args.image):
        print(f"❌ File not found: {args.image}")
        sys.exit(1)

    try:
        if args.raw:
            with open(args.image, "rb") as img_file:
                base64_img = base64.b64encode(img_file.read()).decode('utf-8')
        else:
            with open(args.image, "r", encoding="utf-8") as f:
                base64_img = f.read().strip()

        result = predict_disease(base64_img)
        print(json.dumps(result, indent=2))
    except Exception as e:
        print(json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2))