import gradio as gr
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer
from huggingface_hub import login
import torch
import os  # For env vars

# Login using secret (secure, no hardcode)
login(os.environ["HF_TOKEN"])

# Model setup (loads once on Space startup)
model_id = "agarkovv/CryptoTrader-LM"
base_model_id = "mistralai/Ministral-8B-Instruct-2410"
MAX_LENGTH = 32768
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"  # Use GPU if available (ZeroGPU on HF)

model = AutoPeftModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
model = model.to(DEVICE)
model.eval()

def predict_trading_decision(prompt: str) -> str:
    """Predict daily trading decision (buy, sell, or hold) for BTC or ETH based on news and historical prices.
    
    Args:
        prompt: Input prompt containing cryptocurrency news and historical price data (format: [INST]YOUR PROMPT HERE[/INST]).
    
    Returns:
        Generated trading decision as text (e.g., 'Buy BTC at $62k').
    """
    # Format prompt as required
    formatted_prompt = f"[INST]{prompt}[/INST]"
    
    inputs = tokenizer(
        formatted_prompt, return_tensors="pt", padding=False, max_length=MAX_LENGTH, truncation=True
    )
    inputs = {key: value.to(model.device) for key, value in inputs.items()}

    res = model.generate(
        **inputs,
        use_cache=True,
        max_new_tokens=MAX_LENGTH,
    )
    output = tokenizer.decode(res[0], skip_special_tokens=True)
    return output

# Gradio Interface
demo = gr.Interface(
    fn=predict_trading_decision,
    inputs=gr.Textbox(label="Input Prompt (News + Prices)"),
    outputs=gr.Textbox(label="Trading Decision"),
    title="CryptoTrader-LM MCP Tool",
    description="Predict buy/sell/hold for BTC/ETH."
)

# Launch with MCP support
demo.launch(mcp_server=True)