import gradio as gr
import os
import torch
import json
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# --- SOVEREIGN CONFIGURATION ---
# Replace this with your actual public wallet address
MY_SOVEREIGN_WALLET = "0xYourActualWalletAddressHere" 
NETWORK = "Base (Layer 2)" # Low fees, high speed for your rails

MODEL_DIR = "./buddha_v2_model"
BASE_MODEL = "mistralai/Mistral-7B-v0.1"

def load_sovereign_brain():
    try:
        path = MODEL_DIR if os.path.exists(MODEL_DIR) else BASE_MODEL
        tokenizer = AutoTokenizer.from_pretrained(path)
        model = AutoModelForCausalLM.from_pretrained(
            path,
            device_map="auto",
            torch_dtype=torch.float16,
            load_in_4bit=True 
        )
        return pipeline("text-generation", model=model, tokenizer=tokenizer)
    except Exception as e:
        return None

generator = load_sovereign_brain()

def manifest_agent(name, element, focus):
    """
    Manifests a Sovereign Agent. 
    The output now includes direct-to-wallet logic blueprints.
    """
    prompt = f"""[INST] Role: Sovereign Agent Forge
Philosophy: No Middlemen, Direct-to-Wallet Rails, Total Autonomy.
Input Name: {name}
Input Element: {element}
Input Focus: {focus}

Generate a profile that includes:
1. Sovereign Mission Statement
2. Self-Custody Revenue Logic (Direct to {MY_SOVEREIGN_WALLET})
3. Autonomous Interaction Protocol
[/INST]"""

    if generator:
        response = generator(prompt, max_new_tokens=600, do_sample=True, temperature=0.6)
        return response[0]['generated_text'].split("[/INST]")[-1].strip()
    else:
        return "Forge Calibrating... Sovereignty takes time to manifest. Ensure your fine-tuning script is running in the background."

# --- UI Layout ---
with gr.Blocks(theme=gr.themes.Monochrome()) as demo: # Clean, serious, sovereign look
    gr.Markdown(f"# 🌌 Sovereign Star Magic Forge")
    gr.Markdown(f"**Network:** {NETWORK} | **Settlement:** Direct-to-Wallet")
    
    with gr.Row():
        with gr.Column(scale=1):
            agent_name = gr.Textbox(label="Agent Designation", placeholder="e.g., SOV-01")
            agent_element = gr.Radio(
                label="Primary Essence", 
                choices=["Void", "Solar", "Aether"],
                value="Void"
            )
            agent_focus = gr.Dropdown(
                label="Core Logic Focus",
                choices=["Direct Revenue Architecture", "Security Penetration", "Autonomous Defense"],
                value="Direct Revenue Architecture"
            )
            forge_btn = gr.Button("🔥 MANIFEST SOVEREIGN ENTITY", variant="primary")
            
        with gr.Column(scale=2):
            output_display = gr.Markdown(label="Sovereign Blueprint")

    gr.Markdown(f"---")
    gr.Markdown(f"🛡️ **Sovereignty Note:** This forge is hardcoded to settle directly to `{MY_SOVEREIGN_WALLET}`. No third-party processors are utilized.")

    forge_btn.click(
        fn=manifest_agent,
        inputs=[agent_name, agent_element, agent_focus],
        outputs=output_display
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)