import gradio as gr
import torch
import torch.nn as nn
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
import spaces

# 1. Model Architecture
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)

# 2. Device and Loading Initialization
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

# Load model globally on CPU first
model = SourceCodeAuthorCheck()
model_path = hf_hub_download(repo_id="assix-research/SourceCodeAuthorCheck-SLM-10M", filename="source_code_classifier.pth")
model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=True))
model.eval()

# 3. Inference Logic with ZeroGPU Decorator
@spaces.GPU
def predict_author(code_snippet):
    if not code_snippet or not code_snippet.strip():
        return "Please paste valid code.", "0.0%"
    
    # ZeroGPU dynamically provides CUDA access inside this decorated function
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
        
    inputs = tokenizer(
        code_snippet, 
        return_tensors="pt", 
        truncation=True, 
        padding="max_length", 
        max_length=1024
    ).to(device)
    
    with torch.no_grad():
        if torch.cuda.is_available():
            with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
                logits = model(inputs['input_ids'], inputs['attention_mask'])
        else:
            logits = model(inputs['input_ids'], inputs['attention_mask'])
            
        prob = torch.sigmoid(logits).item()

    score = round(prob * 100, 2)
    verdict = "🤖 AI Generated" if prob > 0.5 else "👨‍💻 Human Written"
    
    # Move model back to CPU to free up ZeroGPU vRAM for other users
    model.to("cpu")
    
    return verdict, f"{score}%"

# 4. Gradio Interface Construction
demo = gr.Interface(
    fn=predict_author,
    inputs=gr.Code(language="python", label="Paste Python Source Code"),
    outputs=[
        gr.Textbox(label="Verdict"),
        gr.Textbox(label="AI Probability Score")
    ],
    title="SourceCodeAuthorCheck SLM (10M)",
    description="Analyze Python snippets to determine if they were written by a human or generated by an AI model.",
    examples=[
        ["def calculate_tax(gross_salary, deduction):\n    return gross_salary - deduction"],
        ["def process_data_stream_0(data_input: list[dict], strict_validation: bool = True) -> dict:\n    if not data_input:\n        return {'status': 'error', 'message': 'Empty stream'}\n    processed_results = []\n    for idx, item in enumerate(data_input):\n        transformed = {k: str(v).strip().lower() for k, v in item.items()}\n        transformed['_internal_id'] = f'gen_id_0_{idx}'\n        processed_results.append(transformed)\n    return {'status': 'success', 'data': processed_results}"]
    ]
)

if __name__ == "__main__":
    demo.launch()