import json
import wandb
import gradio as gr
import os
import tempfile

def sync_to_wandb(json_file, api_key, project_name, run_name, entity):
    """
    Core logic to authenticate with WandB and log HF trainer state data.
    """
    if not json_file:
        return "❌ Please upload a trainer_state.json file."
    if not api_key:
        return "❌ Please provide your WandB API Key."
    if not project_name:
        return "❌ Please provide a project name."

    try:
        # 1. Authenticate
        os.environ["WANDB_API_KEY"] = api_key.strip()
        
        # 2. Load JSON
        with open(json_file.name, 'r') as f:
            state_data = json.load(f)
        
        # 3. Initialize Run
        run = wandb.init(
            project=project_name,
            name=run_name if run_name else "imported-hf-run",
            entity=entity if entity else None,
            reinit=True
        )
        
        run_url = run.get_url()
        status_logs = [f"🚀 Started WandB run: {run.name}", f"🔗 View at: {run_url}"]

        # 4. Log Metadata/Summary
        summary_keys = [
            "best_global_step", "best_metric", "best_model_checkpoint",
            "epoch", "global_step", "max_steps", "num_train_epochs", 
            "total_flos", "trial_name", "trial_params"
        ]
        
        for key in summary_keys:
            if key in state_data:
                wandb.run.summary[key] = state_data[key]

        # 5. Log History
        if "log_history" in state_data:
            history = state_data["log_history"]
            status_logs.append(f"📊 Found {len(history)} log entries. Syncing...")
            
            for entry in history:
                log_payload = entry.copy()
                step = log_payload.get("step")
                
                if step is not None:
                    wandb.log(log_payload, step=int(step))
                else:
                    wandb.log(log_payload)
            
            status_logs.append("✅ Successfully synced all history.")
        else:
            status_logs.append("⚠️ Warning: No 'log_history' found in JSON.")

        # 6. Finish
        wandb.finish()
        status_logs.append("🏁 Sync complete. WandB run finished.")
        
        return "\n".join(status_logs)

    except Exception as e:
        return f"❌ Error: {str(e)}"

# Define the Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🧊 Hugging Face to WandB Syncer
        Upload your `trainer_state.json` file to sync logs from a previous training session to Weights & Biases.
        """
    )
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload trainer_state.json", file_types=[".json"])
            api_key_input = gr.Textbox(
                label="WandB API Key", 
                placeholder="Paste your API key here...", 
                type="password"
            )
            
            with gr.Row():
                project_input = gr.Textbox(label="Project Name", value="hf-import")
                run_name_input = gr.Textbox(label="Run Name", placeholder="e.g., llama-finetuning-v1")
            
            entity_input = gr.Textbox(
                label="Entity (Optional)", 
                placeholder="Team or Username"
            )
            
            sync_btn = gr.Button("Sync to WandB", variant="primary")
            
        with gr.Column():
            output_log = gr.Textbox(
                label="Status / Logs", 
                interactive=False, 
                lines=15
            )

    sync_btn.click(
        fn=sync_to_wandb,
        inputs=[file_input, api_key_input, project_input, run_name_input, entity_input],
        outputs=output_log
    )

    gr.Markdown(
        "--- \n *Note: You can find your API key at [wandb.ai/authorize](https://wandb.ai/authorize).*"
    )

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