import os
import json
import torch
import torch.nn as nn
from torchvision import models

def ensure_files_exist():
    """
    Ensures that all required files exist for the application to run.
    Creates empty or default files if they don't exist.
    """
    print("Checking for required files...")
    
    # Check for model file
    if not os.path.exists('plant_disease_model.pth'):
        print("Model file not found, creating empty file...")
        with open('plant_disease_model.pth', 'w') as f:
            f.write('')
    
    # Check for class names file
    if not os.path.exists('class_names.json'):
        print("Class names file not found, creating default...")
        default_classes = [
            "Apple___Apple_scab", "Apple___Black_rot", "Apple___Cedar_apple_rust", "Apple___healthy",
            "Cherry___healthy", "Cherry___Powdery_mildew", "Corn___Cercospora_leaf_spot", 
            "Corn___Common_rust", "Corn___healthy", "Corn___Northern_Leaf_Blight",
            "Grape___Black_rot", "Grape___Esca_(Black_Measles)", "Grape___healthy", 
            "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)", "Orange___Haunglongbing_(Citrus_greening)",
            "Peach___Bacterial_spot", "Peach___healthy", "Pepper,_bell___Bacterial_spot", 
            "Pepper,_bell___healthy", "Potato___Early_blight", "Potato___healthy", 
            "Potato___Late_blight", "Squash___Powdery_mildew", "Strawberry___healthy",
            "Strawberry___Leaf_scorch", "Tomato___Bacterial_spot", "Tomato___Early_blight", 
            "Tomato___healthy", "Tomato___Late_blight", "Tomato___Leaf_Mold", 
            "Tomato___Septoria_leaf_spot", "Tomato___Spider_mites Two-spotted_spider_mite", 
            "Tomato___Target_Spot", "Tomato___Tomato_mosaic_virus", 
            "Tomato___Tomato_Yellow_Leaf_Curl_Virus"
        ]
        with open('class_names.json', 'w') as f:
            json.dump(default_classes, f)
    
    # Check for treatments CSV
    if not os.path.exists('crop_diseases_treatments.csv'):
        print("Treatments CSV not found, creating it...")
        try:
            import create_treatments_csv
            create_treatments_csv.create_treatments_csv()
        except Exception as e:
            print(f"Error creating treatments CSV: {e}")
    
    # Create example_images directory
    if not os.path.exists('example_images'):
        print("Creating example_images directory...")
        os.makedirs('example_images', exist_ok=True)
    
    print("File check complete.")

def create_minimal_model():
    """
    Creates a minimal model file that can be loaded without errors.
    """
    try:
        print("Creating minimal model file...")
        model = models.resnet18(weights=None)
        num_classes = 38
        model.fc = nn.Linear(model.fc.in_features, num_classes)
        torch.save(model.state_dict(), 'plant_disease_model.pth')
        print(f"Created minimal model file: {os.path.getsize('plant_disease_model.pth')} bytes")
    except Exception as e:
        print(f"Error creating minimal model: {e}")

if __name__ == "__main__":
    ensure_files_exist()
    if os.path.getsize('plant_disease_model.pth') == 0:
        create_minimal_model()
    print("Setup complete!")