import gradio as gr
import pandas as pd
from huggingface_hub import InferenceClient
import os

# --- Configuration ---
# Uses the Secret 'HF_TOKEN' from your Space settings automatically
HF_TOKEN = os.environ.get("HF_TOKEN")

MODELS = {
    "Llama-3-8B-Instruct": "meta-llama/Meta-Llama-3-8B-Instruct",
    "Mistral-7B-Instruct-v0.3": "mistralai/Mistral-7B-Instruct-v0.3",
    "Gemma-7B-It": "google/gemma-7b-it"
}

def load_and_preview_data(file_obj):
    """
    Triggered immediately when a file is uploaded.
    Returns: Dataframe preview, Info String
    """
    if not file_obj:
        return None, "Waiting for file..."
    
    try:
        if file_obj.name.endswith('.csv'):
            df = pd.read_csv(file_obj.name)
        elif file_obj.name.endswith('.json'):
            df = pd.read_json(file_obj.name)
        else:
            return None, "❌ Error: Please upload CSV or JSON."
        
        # Create info string
        info = f"✅ **Loaded Successfully**\n- **Rows:** {len(df)}\n- **Columns:** {', '.join(df.columns)}"
        
        # Return first 5 rows for preview
        return df.head(5), info
        
    except Exception as e:
        return None, f"❌ Error reading file: {str(e)}"

def generate_code(file_obj, model_choice, user_instruction, target_format):
    """
    Generates the Python script using the hidden HF_TOKEN.
    """
    # 1. Security Check
    if not HF_TOKEN:
        return "❌ CRITICAL ERROR: 'HF_TOKEN' is missing in Space Secrets. Go to Settings > Variables and secrets to add it."
    
    if not file_obj:
        return "⚠️ Please upload a file first."

    # 2. Read Data for Context
    try:
        if file_obj.name.endswith('.csv'):
            df = pd.read_csv(file_obj.name)
        else:
            df = pd.read_json(file_obj.name)
    except:
        return "❌ File error."

    # 3. Construct Prompt
    data_sample = df.head(3).to_markdown(index=False)
    columns_info = str(df.dtypes)
    model_id = MODELS[model_choice]
    
    system_prompt = "You are an expert Python Data Engineer. Write ONLY valid Python code. No markdown formatting."
    user_prompt = f"""
    I have a dataset with these columns:
    {columns_info}

    Sample Data:
    {data_sample}

    TASK:
    Write a standalone Python script to convert this dataset into **{target_format}** format for LLM fine-tuning.
    
    USER REQUIREMENTS:
    {user_instruction}
    
    OUTPUT:
    - Use 'pandas' library.
    - Handle missing values if necessary.
    - Save output to 'ready_for_finetune.jsonl'.
    - Return ONLY the code.
    """
    
    # 4. Call Model
    try:
        client = InferenceClient(model=model_id, token=HF_TOKEN)
        response = client.chat_completion(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            max_tokens=1500,
            temperature=0.1
        )
        code = response.choices[0].message.content
        return code.replace("```python", "").replace("```", "").strip()
        
    except Exception as e:
        return f"❌ Inference Error: {str(e)}"

# --- Advanced UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🧬 Intelligent Dataset Prep
        ### Auto-generate cleaning scripts using Llama 3 & Mistral
        """
    )
    
    # Section 1: Data Viewer (Full Width)
    with gr.Row():
        with gr.Column(scale=1):
            file_upload = gr.File(label="📂 Upload Raw Dataset (CSV/JSON)", file_types=[".csv", ".json"])
        with gr.Column(scale=2):
            file_info = gr.Markdown("Waiting for upload...")
            data_preview = gr.DataFrame(label="🔍 Data Preview (First 5 Rows)", interactive=False)

    gr.Markdown("---")

    # Section 2: Controls & Output
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Configuration")
            
            model_select = gr.Dropdown(
                choices=list(MODELS.keys()), 
                value="Llama-3-8B-Instruct", 
                label="Select AI Model"
            )
            
            format_select = gr.Dropdown(
                choices=["Alpaca (Instruction)", "ShareGPT (Chat)", "HuggingFace Format"], 
                value="Alpaca (Instruction)", 
                label="Target Format"
            )
            
            instructions = gr.Textbox(
                label="Transformation Instructions", 
                value="Combine 'title' and 'summary' into 'instruction'. Use 'response' as output. Drop nulls.",
                lines=4,
                placeholder="Describe how to map your columns..."
            )
            
            btn_run = gr.Button("⚡ Generate Script", variant="primary", size="lg")
            
        with gr.Column(scale=1):
            gr.Markdown("### 🐍 Generated Python Code")
            code_out = gr.Code(language="python", label="Script", lines=20)

    # --- Interaction Logic ---
    # 1. When file is uploaded, update the preview and info text
    file_upload.change(
        load_and_preview_data,
        inputs=[file_upload],
        outputs=[data_preview, file_info]
    )
    
    # 2. When button clicked, generate code
    btn_run.click(
        generate_code,
        inputs=[file_upload, model_select, instructions, format_select],
        outputs=[code_out]
    )

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