import json
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr

with open("config.json") as f:
    CFG = json.load(f)

with open("tokenizer.json") as f:
    tokenizer = json.load(f)

with open("threshold.txt") as f:
    THRESHOLD = float(f.read().strip())

CHAR2ID = tokenizer["char2id"]

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)

        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, d_model, 2)
            * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return self.dropout(x + self.pe[:, :x.size(1)])

class PhishingTransformer(nn.Module):

    def __init__(self):
        super().__init__()

        self.embed = nn.Embedding(
            CFG["vocab_size"],
            CFG["embed_dim"],
            padding_idx=CFG["pad_token"]
        )

        self.pos = PositionalEncoding(
            CFG["embed_dim"],
            CFG["max_len"],
            CFG["dropout"]
        )

        layer = nn.TransformerEncoderLayer(
            d_model=CFG["embed_dim"],
            nhead=CFG["num_heads"],
            dim_feedforward=CFG["ff_dim"],
            dropout=CFG["dropout"],
            batch_first=True,
            activation="gelu",
            norm_first=True
        )

        self.encoder = nn.TransformerEncoder(
            layer,
            num_layers=CFG["num_layers"]
        )

        self.norm = nn.LayerNorm(CFG["embed_dim"])

        self.head = nn.Sequential(
            nn.Linear(CFG["embed_dim"], CFG["embed_dim"] // 2),
            nn.GELU(),
            nn.Dropout(CFG["dropout"]),
            nn.Linear(CFG["embed_dim"] // 2, 2)
        )

    def forward(self, x):

        mask = (x == CFG["pad_token"])

        x = self.embed(x)

        x = self.pos(x)

        x = self.encoder(
            x,
            src_key_padding_mask=mask
        )

        valid = (~mask).unsqueeze(-1).float()

        x = (x * valid).sum(1) / valid.sum(1).clamp(min=1e-9)

        x = self.norm(x)

        return self.head(x)

model = PhishingTransformer()

model.load_state_dict(
    torch.load(
        "phishing_detector.pth",
        map_location="cpu"
    )
)

model.eval()

def predict(url):

    ids = [
        CHAR2ID.get(c, 1)
        for c in url[:CFG["max_len"]]
    ]

    while len(ids) < CFG["max_len"]:
        ids.append(CFG["pad_token"])

    x = torch.tensor([ids])

    with torch.no_grad():
        logits = model(x)
        probs = F.softmax(logits, dim=-1)[0]

    score = float(probs[1])

    return {
        "phishing": score >= THRESHOLD,
        "score": score
    }

demo = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="json",
    title="Phishing Detector"
)

demo.launch()