import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModel
import torch.nn as nn

class CodeClassifier(nn.Module):
    def __init__(self, model_name="microsoft/codebert-base", num_labels=2):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(model_name)
        self.classifier = nn.Linear(768, num_labels)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        return self.classifier(outputs.last_hidden_state[:, 0, :])

def load_model():
    model = CodeClassifier(num_labels=2)
    
    try:
        from huggingface_hub import hf_hub_download
        model_path = hf_hub_download(
            repo_id="KrishnaKarthik/ai-code-detector",
            filename="ai_code_detector.pth"
        )
        model.load_state_dict(torch.load(model_path, map_location="cpu"))
        print("✅ Loaded your trained AI model!")
    except Exception as e:
        print(f"❌ Error: {str(e)}")
    
    model.eval()
    return model

model = load_model()
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")

def analyze_code(code):
    if not code.strip():
        return "Please enter some code", ""
    
    inputs = tokenizer(code, return_tensors="pt", truncation=True, max_length=512)
    
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs, dim=1)
        human_prob = probabilities[0][0].item()
        ai_prob = probabilities[0][1].item()
    
    human_percent = f"{human_prob:.1%}"
    ai_percent = f"{ai_prob:.1%}"
    
    if human_prob > 0.7:
        verdict = f"✅ Likely Human-written ({human_prob:.1%} confidence)"
    elif ai_prob > 0.7:
        verdict = f"🤖 Likely AI-generated ({ai_prob:.1%} confidence)"
    else:
        verdict = f"⚠️ Uncertain - could be mixed"
    
    return human_percent, ai_percent, verdict

# Create Gradio interface
with gr.Blocks(title="AI Code Detector") as demo:
    gr.Markdown("# 🔍 AI-Generated Code Detector")
    gr.Markdown("Paste any code to check if it's AI-generated or human-written")
    
    with gr.Row():
        code_input = gr.Textbox(
            label="Paste your code here",
            placeholder="def hello_world():\n    print('Hello, World!')",
            lines=10
        )
    
    analyze_btn = gr.Button("Analyze Code", variant="primary")
    
    with gr.Row():
        human_output = gr.Textbox(label="Human-written Probability")
        ai_output = gr.Textbox(label="AI-generated Probability")
    
    verdict_output = gr.Textbox(label="Verdict")
    
    analyze_btn.click(
        fn=analyze_code,
        inputs=code_input,
        outputs=[human_output, ai_output, verdict_output]
    )

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