import spaces
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "PraneetNS/codesentinel-full"
TOKENIZER_ID = "PraneetNS/codesentinel-adapter"

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    TOKENIZER_ID,
    trust_remote_code=True,
)

SYSTEM_PROMPT = """You are CodeSentinel.

You are an expert AI software engineer.

Capabilities:
- Detect bugs
- Explain code
- Fix code
- Secure code
- Refactor code
- Explain algorithms
- Generate production-quality software

Never generate malware or unsafe code.
"""

model = None


def load_model():
    global model

    if model is None:
        print("Loading model...")

        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            dtype=torch.float16,
            device_map="auto",
            trust_remote_code=True,
        )

        model.eval()

    return model


@spaces.GPU
def generate(message, history):
    model = load_model()

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        }
    ]

    # Convert Gradio history
    if history:
        for msg in history:
            role = msg.get("role")

            content = msg.get("content", "")

        # Gradio 6 sometimes stores content as a list
            if isinstance(content, list):
                text = ""

                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") == "text":
                            text += part.get("text", "")
                    elif isinstance(part, str):
                        text += part

                content = text

            messages.append(
                {
                "role": role,
                "content": content,
            }
        )
    messages.append(
    {
        "role": "user",
        "content": message,
    }
)

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )

    inputs = {
        k: v.to(model.device)
        for k, v in inputs.items()
    }

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1024,
            temperature=0.3,
            top_p=0.95,
            do_sample=True,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]

    response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    )

    return response.strip()


demo = gr.ChatInterface(
    fn=generate,
    title="🛡️ CodeSentinel",
    description="AI-powered Secure Coding Assistant",
    examples=[
        "Find the bug in this Python code.",
        "Explain this C++ function.",
        "Optimize this SQL query.",
        "Write a secure FastAPI login API.",
        "Review this Java code for security vulnerabilities.",
    ],
)


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
    )