import torch
import torch.nn as nn
from torchvision import models
import os
import urllib.request
import shutil
import ssl  # <--- Add this import!

# --- AUTO-DOWNLOAD WEIGHTS SETUP ---
HOME_DIR = os.path.expanduser("~")
WEIGHTS_DIR = os.path.join(HOME_DIR, ".graphvision_weights")
os.makedirs(WEIGHTS_DIR, exist_ok=True)

HF_BASE_URL = "https://huggingface.co/ShadowGard3n/graphvision/resolve/main/"

WEIGHT_URLS = {
    "graph_classifier.pth": f"{HF_BASE_URL}graph_classifier_real.pth",
    "pie_regressor_stable.pth": f"{HF_BASE_URL}pie_regressor_stable.pth",
    "hbar_regressor_stable.pth": f"{HF_BASE_URL}hbar_regressor_stable.pth",
    "vbar_regressor.pth": f"{HF_BASE_URL}vbar_regressor.pth"
}

def download_weight(filename):
    """Downloads the weight file from Hugging Face if it doesn't already exist."""
    filepath = os.path.join(WEIGHTS_DIR, filename)
    if not os.path.exists(filepath):
        print(f"Downloading {filename} from Hugging Face (This only happens once)...")
        req = urllib.request.Request(WEIGHT_URLS[filename], headers={'User-Agent': 'Mozilla/5.0'})
        
        try:
            # 1. Try the standard secure download
            with urllib.request.urlopen(req) as response, open(filepath, 'wb') as out_file:
                shutil.copyfileobj(response, out_file)
        except urllib.error.URLError as e:
            # 2. If it's the classic Mac SSL error, use a fallback context to force the download
            if "CERTIFICATE_VERIFY_FAILED" in str(e.reason):
                print("Mac SSL bug detected! Using unverified context to securely force download...")
                context = ssl._create_unverified_context()
                with urllib.request.urlopen(req, context=context) as response, open(filepath, 'wb') as out_file:
                    shutil.copyfileobj(response, out_file)
            else:
                raise e # If it's a different error (like no internet), raise it normally
                
    return filepath

# --- 1. MODEL ARCHITECTURES ---
class PieRegressor(nn.Module):
    def __init__(self):
        super(PieRegressor, self).__init__()
        self.backbone = models.resnet18(weights=None)
        num_ftrs = self.backbone.fc.in_features
        self.backbone.fc = nn.Linear(num_ftrs, 10)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        return self.sigmoid(self.backbone(x))

class BarRegressor(nn.Module):
    def __init__(self):
        super(BarRegressor, self).__init__()
        self.backbone = models.resnet18(weights=None)
        num_ftrs = self.backbone.fc.in_features
        self.backbone.fc = nn.Linear(num_ftrs, 10)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        return self.sigmoid(self.backbone(x))

# --- 2. SECURE LOADING FUNCTIONS ---
def load_classifier(device):
    # FIX: Use the download_weight function instead of os.path.join
    path = download_weight("graph_classifier.pth")
    
    model = models.resnet18(weights=None)
    num_ftrs = model.fc.in_features
    # Set to 5 classes to perfectly match the saved checkpoint
    model.fc = nn.Linear(num_ftrs, 5) 
    
    model.load_state_dict(torch.load(path, map_location=device))
    return model.to(device).eval()

def load_pie_model(device):
    # FIX: Use the download_weight function
    path = download_weight("pie_regressor_stable.pth")
    model = PieRegressor()
    model.load_state_dict(torch.load(path, map_location=device))
    return model.to(device).eval()

def load_hbar_model(device):
    # FIX: Use the download_weight function
    path = download_weight("hbar_regressor_stable.pth")
    model = BarRegressor()
    model.load_state_dict(torch.load(path, map_location=device))
    return model.to(device).eval()

def load_vbar_model(device):
    # FIX: Use the download_weight function
    path = download_weight("vbar_regressor.pth") 
    model = BarRegressor()
    model.load_state_dict(torch.load(path, map_location=device))
    return model.to(device).eval()