import sys
import os
import torch
import torch.nn as nn
from transformers import AutoTokenizer

class SourceCodeAuthorCheck(nn.Module):
    def __init__(self, vocab_size=50257, d_model=128, nhead=8, num_layers=4, dim_feedforward=512):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = nn.Parameter(torch.zeros(1, 1024, d_model))
        
        encoder_layers = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layers, num_layers=num_layers)
        self.fc = nn.Linear(d_model, 1)

    def forward(self, input_ids, attention_mask):
        seq_len = input_ids.size(1)
        x = self.embedding(input_ids) + self.pos_encoder[:, :seq_len, :]
        
        src_key_padding_mask = ~attention_mask.bool()
        x = self.transformer(x, src_key_padding_mask=src_key_padding_mask)
        
        mask_expanded = attention_mask.unsqueeze(-1).float()
        sum_embeddings = torch.sum(x * mask_expanded, 1)
        sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
        pooled = sum_embeddings / sum_mask
        
        return self.fc(pooled)

def predict(code_snippet, model, tokenizer, device):
    inputs = tokenizer(
        code_snippet, 
        return_tensors="pt", 
        truncation=True, 
        padding="max_length", 
        max_length=1024
    ).to(device)
    
    with torch.no_grad():
        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
            logits = model(inputs['input_ids'], inputs['attention_mask'])
            prob = torch.sigmoid(logits).item()

    return prob

def main():
    if len(sys.argv) < 2:
        print("Usage: python 4-inference_file.py <path_to_python_file>")
        sys.exit(1)

    target_file = sys.argv[1]
    
    if not os.path.exists(target_file):
        print(f"Error: File '{target_file}' not found.")
        sys.exit(1)

    with open(target_file, 'r', encoding='utf-8', errors='ignore') as f:
        code_content = f.read()

    if not code_content.strip():
        print(f"Error: File '{target_file}' is empty.")
        sys.exit(1)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    model = SourceCodeAuthorCheck().to(device)
    model.load_state_dict(torch.load("source_code_classifier.pth", map_location=device, weights_only=True))
    model.eval()

    print(f"\n--- Testing File: {target_file} ---")
    prob = predict(code_content, model, tokenizer, device)
    score = round(prob * 100, 2)
    
    verdict = "AI Generated" if prob > 0.5 else "Human Written"
    print(f"Verdict: {verdict} (AI Probability: {score}%)")
    print(f"Preview: {code_content[:150].strip()}...\n")

if __name__ == "__main__":
    main()