import os
import spaces
import gradio as gr
from pathlib import Path

from extractor import load_model, extract

print("MemTotal:", open("/proc/meminfo").readline().strip())
print("cgroup limit:", open("/sys/fs/cgroup/memory.max").read().strip()
      if os.path.exists("/sys/fs/cgroup/memory.max") else "n/a")

llm = load_model()


@spaces.GPU
def gpu_placeholder():
    return None


def analyze(uploaded_file):
    if uploaded_file is None:
        return "Please upload a .py file."
    try:
        path = Path(uploaded_file)
        result = extract(llm, path.name, path.read_text(encoding="utf-8"))
        return result.model_dump_json(indent=2)
    except Exception as e:
        return f"Error: {e}"


demo = gr.Interface(
    fn=analyze,
    inputs=gr.File(label="Upload a Python file", file_types=[".py"]),
    outputs=gr.Code(label="Extraction result", language="json"),
    title="Function & Call Extractor",
)

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