from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import os
import socket
from urllib.parse import urlparse
from models.gnn import PhishingGNN_Model
from pipeline.graph_engine import TopologicalGraphEngine
from config import *
import uvicorn

app = FastAPI(title="Defender V5 Sovereign Threat Analysis API")

# Define node dimension configurations
in_channels_dict = {
    'ip': 16,
    'domain': 32,
    'asn': 8,
    'cert': 16
}

# Load model directly into state memory matching training dimension maps
model = PhishingGNN_Model(
    metadata=GRAPH_METADATA,
    in_channels_dict=in_channels_dict,
    hidden_channels=HIDDEN_CHANNELS,
    num_heads=NUM_HEADS,
    num_layers=NUM_LAYERS,
    dropout_rate=0.0
)

if os.path.exists(MODEL_SAVE_PATH):
    model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location='cpu', weights_only=True))
model.eval()

# Payload model expects a URL
class URLPayload(BaseModel):
    url: str

@app.post("/analyze")
async def analyze_url(payload: URLPayload):
    try:
        # 1. Parse the URL to extract the domain
        parsed_url = urlparse(payload.url)
        domain = parsed_url.netloc or parsed_url.path.split('/')[0]
        
        if ':' in domain:
            domain = domain.split(':')[0]

        if not domain:
            raise HTTPException(status_code=400, detail="Invalid URL format.")

        # 2. Resolve the domain to an IP address
        try:
            ip = socket.gethostbyname(domain)
        except socket.gaierror:
            raise HTTPException(status_code=400, detail=f"DNS Resolution failed for domain: {domain}")

        # 3. Build telemetry dictionary
        telemetry_log = {
            "ip": ip,
            "domain": domain,
            "asn": None 
        }

        # 4. Pass through Graph Engine
        engine = TopologicalGraphEngine()
        x_dict, edge_index_dict = engine.extract_and_build([telemetry_log])
        
        with torch.no_grad():
            raw_scores = model(x_dict, edge_index_dict)
            threat_probability = float(raw_scores.max().item())
            
        return {
            "input_url": payload.url,
            "resolved_domain": domain,
            "resolved_ip": ip,
            "structural_anomaly_score": round(threat_probability, 5),
            "remediation_verdict": "ISOLATE_ROUTING" if threat_probability > 0.70 else "ALLOW"
        }
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Graph Engine Exception: {str(e)}")

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8005)