import json
import gradio as gr
from server.environment import RevOpsEnvironment
from models import RevOpsAction, ActionType

env = RevOpsEnvironment()
env.reset(task_id="task_easy")

def format_lead(obs):
    if not obs or not hasattr(obs, 'lead'):
         return "No lead data."
    lead = obs.lead
    return f"""### Current Lead
**ID**: {lead.id}
**Name**: {lead.name}
**Title**: {lead.title}
**Email**: {lead.email}
**Company**: {lead.company}
**Source**: {lead.source}
**Message**: {lead.message}
**Score**: {lead.score}"""

def format_enrichment(obs):
    if not obs or not obs.enrichment.enriched:
         return "Not Enriched Yet."
    e = obs.enrichment
    return f"""### Firmographic Data
**Company Size**: {e.company_size}
**Revenue**: {e.revenue}
**Industry**: {e.industry}
**Region**: {e.region}
**Tier**: {e.tier}"""

def format_crm(obs):
    if not obs or not obs.crm.checked:
         return "CRM Not Checked."
    accounts = obs.crm.existing_accounts
    opps = obs.crm.opportunities
    
    res = "### CRM Match Results\n"
    if not accounts and not opps:
         return res + "No existing CRM records found."
    
    if accounts:
         res += "#### Existing Accounts:\n"
         for a in accounts:
              res += f"- {a.get('name')} (ID: {a.get('account_id')})\n"
    
    if opps:
         res += "\n#### Opportunities:\n"
         for o in opps:
              res += f"- Opp ID: {o.opportunity_id} | Status: {o.status} | Amount: ${o.amount} | Assigned Rep: {o.assigned_rep_id}\n"
    return res

def refresh_ui():
    st = env.get_full_state()
    obs = env._get_observation()
    done = obs.done if hasattr(obs, 'done') else False
    
    feedback = st.get("last_feedback", "")
    score = st.get("grader_score")
    
    status_text = f"**Status**: {'DONE' if done else 'IN PROGRESS'}\n\n**Accumulated Reward**: {st.get('accumulated_reward', 0.0):.2f}"
    if score is not None:
         status_text += f"\n\n**Final Grader Score**: {score}"
         
    return (
        format_lead(obs),
        format_enrichment(obs),
        format_crm(obs),
        feedback,
        status_text
    )

def reset_task(task_id):
    env.reset(task_id=task_id)
    return refresh_ui()

def perform_action(action_type, score, rep_id, account_id, opportunity_id, disqualification_reason):
    st = env.get_full_state()
    if st.get("current_lead_index", 0) >= len(st.get("leads", [])):
         return refresh_ui() # Already done
         
    # Convert string action back to enum safely
    action_val = ActionType(action_type)
    
    action = RevOpsAction(
        action_type=action_val,
        score=int(score) if score else None,
        rep_id=rep_id if rep_id else None,
        account_id=account_id if account_id else None,
        opportunity_id=opportunity_id if opportunity_id else None,
        disqualification_reason=disqualification_reason if disqualification_reason else None
    )
    
    env.step(action)
    return refresh_ui()

# --- Gradio UI Layout ---
with gr.Blocks(title="RevOps CRM Agent Desktop", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🏢 RevOps-Agent Live Dashboard")
    gr.Markdown("Interact with the OpenEnv OpenEnv backend as a human operator, or watch your agent's live states.")
    
    with gr.Row():
        with gr.Column(scale=2):
            task_dropdown = gr.Dropdown(choices=["task_easy", "task_medium", "task_hard"], value="task_easy", label="Select Simulation Scenario")
            reset_btn = gr.Button("Reset Environment", variant="primary")
            
        with gr.Column(scale=1):
            status_box = gr.Markdown(value="**Status**: IN PROGRESS\n**Accumulated Reward**: 0.00")
            
    with gr.Row():
        with gr.Column(scale=1):
            lead_box = gr.Markdown()
        with gr.Column(scale=1):
            enrichment_box = gr.Markdown()
        with gr.Column(scale=1):
            crm_box = gr.Markdown()
            
    with gr.Row():
        feedback_box = gr.Textbox(label="Last Action Feedback", interactive=False)
            
    gr.Markdown("---")
    gr.Markdown("### ⚡ Action Console")
    with gr.Row():
        with gr.Column(scale=1):
            action_dropdown = gr.Dropdown(choices=[e.value for e in ActionType], value="enrich_lead", label="Action Type")
            btn_submit = gr.Button("Execute Action", variant="primary")
            
        with gr.Column(scale=2):
            score_input = gr.Number(label="Score (0-100) - For update_lead_score")
            rep_id_input = gr.Textbox(label="Rep ID - For route_to_rep (e.g., rep_amer_mm, rep_emea_ent)")
            account_id_input = gr.Textbox(label="Account ID - For merge_with_account")
            opp_id_input = gr.Textbox(label="Opportunity ID - For flag_reengagement")
            disqualify_input = gr.Textbox(label="Disqualification Reason - For disqualify")

    # Wire events
    outputs = [lead_box, enrichment_box, crm_box, feedback_box, status_box]
    
    demo.load(refresh_ui, inputs=None, outputs=outputs)
    
    reset_btn.click(reset_task, inputs=[task_dropdown], outputs=outputs)
    
    btn_submit.click(
        perform_action, 
        inputs=[action_dropdown, score_input, rep_id_input, account_id_input, opp_id_input, disqualify_input],
        outputs=outputs
    )

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