import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
import json

# Use a smaller 7B model for demo
MODEL_REPO = "TheBloke/WizardLM-7B-V1.0-GGUF"
MODEL_FILE = "wizardlm-7b-v1.0.Q4_K_M.gguf"

# Download model automatically
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)

# Load quantized GGUF model
llm = Llama(
    model_path=model_path,
    n_ctx=2048,   # reduce context for speed
    n_threads=4   # keep light for demo
)

# Function to optimize loads
def optimize_loads(input_json: str):
    try:
        data = json.loads(input_json)
    except Exception as e:
        return f"Invalid JSON: {e}"

    prompt = f"""
You are a truck load optimizer. 
Here is the truck and package data in JSON:
{json.dumps(data, indent=2)}

Task:
- Place packages into front, middle, rear sections.
- Distribute left/right to maintain balance.
- Respect weight and volume limits.
- Return ONLY valid JSON output with this structure:

{{
  "front": [],
  "middle": [],
  "rear": [],
  "left_side_weight": 0,
  "right_side_weight": 0
}}
    """

    response = llm(prompt, max_tokens=256, stop=["}"])
    text = response["choices"][0]["text"]

    try:
        return json.dumps(json.loads(text + "}"), indent=2)
    except:
        return text.strip()

# Gradio UI
iface = gr.Interface(
    fn=optimize_loads,
    inputs=gr.Textbox(label="Truck & Load Data (JSON)", lines=15, placeholder="Paste JSON here"),
    outputs=gr.Textbox(label="Optimized Load Plan (JSON)", lines=15),
    title="Truck Load Optimizer (Demo - 7B Model)",
    description="Paste truck & package JSON → get optimal load plan.",
    examples=[["sample_input.json"]]
)

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