import os
import sys
import time
import asyncio
import torch
import torch.nn.functional as F
import timm
from torchvision import transforms
from PIL import Image
from playwright.async_api import async_playwright

class ProductionAnalyzer:
    def __init__(self, model_path="models/production_resnet_ema.pth", model_name="resnet18", num_classes=2):
        """
        Initializes the ResNet18 production model on available hardware.
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"🖥️ Initializing backend hardware: {self.device}")
        
        # 1. Initialize the ResNet18 architecture base
        print(f"🏗️ Building network architecture: {model_name}...")
        self.model = timm.create_model(model_name, pretrained=False, num_classes=num_classes)
        
        # 2. Safety check for the weights file
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"❌ Could not find weight file at: {model_path}\n"
                                    f"Please ensure it is placed inside the 'models' folder.")
            
        print(f"📥 Loading ResNet18 weights from {model_path}...")
        state_dict = torch.load(model_path, map_location=self.device, weights_only=True)
        
        # 3. 🛠️ FIXED: Strip wrapper prefixes and safely drop training metadata keys
        clean_state_dict = {}
        for key, value in state_dict.items():
            if key == "n_averaged":
                continue # Skip the training counter metadata so PyTorch doesn't throw an error
                
            clean_key = key.replace('module.', '')
            clean_state_dict[clean_key] = value
            
        self.model.load_state_dict(clean_state_dict)
        
        # 4. Lock the model for evaluation mode
        self.model.eval()
        self.model.to(self.device)
        print("✅ ResNet18 Engine successfully loaded and locked for inference.")
        
        # 5. Define standard normalizations expected by ResNet18
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
        
        # Alphabetical class mapping (0: Legit, 1: Phishing)
        self.class_names = ["Legit", "Phishing"]

    def analyze_image(self, image_path):
        """
        Feeds the captured screenshot into the ResNet18 neural network.
        """
        try:
            image = Image.open(image_path).convert('RGB')
            input_tensor = self.transform(image).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                logits = self.model(input_tensor)
                probabilities = F.softmax(logits[0], dim=0)
                confidence, predicted_idx = torch.max(probabilities, 0)
                
            return {
                "prediction": self.class_names[predicted_idx.item()],
                "confidence": f"{confidence.item() * 100:.2f}%"
            }
        except Exception as e:
            return {"error": f"Model inference failed: {str(e)}"}


async def capture_screenshot(url, output_path="temp_inference.png"):
    """
    Launches a headless browser to safely capture a screenshot of the live URL.
    """
    if not url.startswith('http://') and not url.startswith('https://'):
        url = 'https://' + url
        
    print(f"🌐 Navigating to: {url} ...")
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        context = await browser.new_context(
            viewport={'width': 1280, 'height': 720},
            ignore_https_errors=True, # Bypasses broken SSL certifications on malicious sites
            accept_downloads=False
        )
        
        page = await context.new_page()
        try:
            # 12-second timeout to handle slow/malicious servers
            await page.goto(url, timeout=12000, wait_until='domcontentloaded')
            await asyncio.sleep(1) # Brief pause to let visual components load completely
            await page.screenshot(path=output_path)
            return output_path
        except Exception as e:
            print(f"❌ Failed to reach or capture the website: {e}")
            return None
        finally:
            await page.close()
            await browser.close()


async def main():
    print("=" * 45)
    print("🛡️ CHIMERA 2.0 LIVE URL DETECTOR ENGINE")
    print("=" * 45)
    
    try:
        analyzer = ProductionAnalyzer()
    except Exception as e:
        print(e)
        return

    TEMP_IMG = "temp_inference.png"
    
    try:
        while True:
            print("\n" + "-" * 45)
            url_input = input("🔗 Enter URL to inspect (or type 'exit' to quit): ").strip()
            
            if url_input.lower() == 'exit':
                print("Shutting down engine...")
                break
            if not url_input:
                continue
                
            start_time = time.time()
            
            # Step 1: Capture screenshot via Playwright
            screenshot_file = await capture_screenshot(url_input, TEMP_IMG)
            
            if screenshot_file and os.path.exists(screenshot_file):
                # Step 2: Pass screenshot into ResNet18
                print("🔍 Running Deep Learning Visual Inspection...")
                result = analyzer.analyze_image(screenshot_file)
                
                total_time = time.time() - start_time
                
                # Step 3: Output results safely
                print("\n" + "=" * 35)
                print("📊 LIVE DETECTION REPORT")
                print("=" * 35)
                if "error" in result:
                    print(f"Result: {result['error']}")
                else:
                    status_prefix = "🚨 ALERT!!" if result['prediction'] == "Phishing" else "✅ CLEAR:"
                    print(f"Verdict    : {status_prefix} {result['prediction']}")
                    print(f"Confidence : {result['confidence']}")
                    print(f"Total Time : {total_time:.2f} seconds")
                print("=" * 35)
                
                if os.path.exists(TEMP_IMG):
                    os.remove(TEMP_IMG)
            else:
                print("❌ Inspection aborted. Visual fingerprint could not be gathered.")
                
    except KeyboardInterrupt:
        print("\nExiting execution gracefully...")
    finally:
        if os.path.exists(TEMP_IMG):
            os.remove(TEMP_IMG)

if __name__ == "__main__":
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
        
    asyncio.run(main())