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

model_id = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
hf_token = os.environ.get("HF_TOKEN")

print("Loading model securely...")
tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    dtype=torch.float32, 
    token=hf_token
)

def solve(problem_text):
    if not problem_text or len(problem_text) < 10:
        return "// Error: Problem text too short."

    prompt = f"Problem:\n{problem_text}\n\nOptimal and correct Python3 code solution:\n"
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_new_tokens=700, 
            do_sample=False, 
            pad_token_id=tokenizer.eos_token_id
        )
    
    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 1. Extract the raw Python code block
    try:
        code_block = full_text.split("```python")[1]
        pure_code = code_block.split("```")[0]
    except IndexError:
        if "Python code solution:\n" in full_text:
            pure_code = full_text.split("Python code solution:\n")[1]
        else:
            pure_code = full_text

    # 2. THE SCRUBBER: Delete the comments using Python, not the AI
    cleaned_lines = []
    for line in pure_code.split('\n'):
        # This splits the line at the '#' symbol and only keeps the code on the left side
        no_comment_line = line.split('#')[0].rstrip()
        
        # Only add the line if it's not completely empty
        if no_comment_line.strip():
            cleaned_lines.append(no_comment_line)

    final_code = "\n".join(cleaned_lines).strip()
    
    # Safety net: If the scrubber accidentally deleted everything, return the original
    if not final_code:
        return pure_code.strip()
        
    return final_code

demo = gr.Interface(fn=solve, inputs="text", outputs="text", api_name="predict")
demo.launch()