import gradio as gr
import json
import random
from typing import Dict, List, Any, Tuple

# Simulate OMC Talent Market data
TALENT_DATABASE = {
    "research_analyst": {
        "name": "ResearchAnalyst",
        "skills": ["web_search", "summarization", "fact_checking"],
        "tools": ["browser", "search_api", "document_reader"],
        "backend": "gpt-4-class",
        "cost_per_task": 0.02,
        "reliability": 0.95
    },
    "code_generator": {
        "name": "CodeGenerator", 
        "skills": ["python", "javascript", "api_integration"],
        "tools": ["code_interpreter", "git", "shell"],
        "backend": "claude-class",
        "cost_per_task": 0.03,
        "reliability": 0.92
    },
    "data_processor": {
        "name": "DataProcessor",
        "skills": ["pandas", "sql", "visualization"],
        "tools": ["sqlite", "matplotlib", "csv_parser"],
        "backend": "local-llm",
        "cost_per_task": 0.01,
        "reliability": 0.88
    },
    "review_auditor": {
        "name": "ReviewAuditor",
        "skills": ["quality_check", "error_detection", "consistency_verification"],
        "tools": ["diff_tool", "test_runner", "linter"],
        "backend": "gpt-4-class",
        "cost_per_task": 0.025,
        "reliability": 0.97
    },
    "planner_orchestrator": {
        "name": "PlannerOrchestrator",
        "skills": ["task_decomposition", "dependency_analysis", "scheduling"],
        "tools": ["graph_builder", "timeline_tracker"],
        "backend": "o1-class",
        "cost_per_task": 0.05,
        "reliability": 0.93
    }
}

TASK_TEMPLATES = [
    {
        "id": "t1",
        "description": "Research competitor pricing and generate comparison report",
        "required_skills": ["web_search", "summarization", "pandas", "visualization"],
        "complexity": "medium",
        "budget": 0.10
    },
    {
        "id": "t2", 
        "description": "Build Python API client for REST service with error handling",
        "required_skills": ["python", "api_integration", "quality_check"],
        "complexity": "high",
        "budget": 0.12
    },
    {
        "id": "t3",
        "description": "Analyze CSV sales data and create trend dashboard",
        "required_skills": ["pandas", "sql", "visualization"],
        "complexity": "low", 
        "budget": 0.06
    },
    {
        "id": "t4",
        "description": "Review code for security vulnerabilities in authentication module",
        "required_skills": ["quality_check", "error_detection", "python"],
        "complexity": "high",
        "budget": 0.15
    }
]

def calculate_skill_match(talent_skills: List[str], required: List[str]) -> float:
    if not required:
        return 1.0
    matches = sum(1 for skill in required if skill in talent_skills)
    return matches / len(required)

def recruit_talents_for_task(task_idx: int) -> Tuple[str, str]:
    if task_idx < 0 or task_idx >= len(TASK_TEMPLATES):
        return "Invalid task selection", "{}"
    
    task = TASK_TEMPLATES[task_idx]
    required_skills = set(task["required_skills"])
    budget = task["budget"]
    
    recruited = []
    covered_skills = set()
    total_cost = 0.0
    
    steps_log = [f"🎯 TASK: {task['description']}",
                 f"📋 Required skills: {', '.join(required_skills)}",
                 f"💰 Budget: ${budget:.3f}",
                 "═" * 50]
    
    while covered_skills != required_skills and total_cost < budget:
        best_talent = None
        best_value = 0
        
        for talent_id, talent in TALENT_DATABASE.items():
            if talent_id in [r["id"] for r in recruited]:
                continue
                
            new_coverage = required_skills - covered_skills
            talent_skills = set(talent["skills"])
            covered_by_this = len(new_coverage & talent_skills)
            
            if covered_by_this == 0:
                continue
                
            value = (covered_by_this / talent["cost_per_task"]) * talent["reliability"]
            
            if value > best_value and total_cost + talent["cost_per_task"] <= budget:
                best_value = value
                best_talent = talent_id
        
        if best_talent is None:
            break
            
        talent = TALENT_DATABASE[best_talent]
        recruited.append({
            "id": best_talent,
            "name": talent["name"],
            "cost": talent["cost_per_task"],
            "skills_added": list(set(talent["skills"]) & required_skills - covered_skills)
        })
        
        for skill in talent["skills"]:
            covered_skills.add(skill)
            
        total_cost += talent["cost_per_task"]
        
        steps_log.append(f"✅ RECRUITED: {talent['name']}")
        steps_log.append(f"   💵 Cost: ${talent['cost_per_task']:.3f}")
        steps_log.append(f"   🛠️  Skills covered: {', '.join(talent['skills'])}")
        steps_log.append(f"   📊 Coverage: {len(covered_skills & required_skills)}/{len(required_skills)} skills")
        steps_log.append("")
    
    uncovered = required_skills - covered_skills
    if uncovered:
        steps_log.append(f"⚠️  WARNING: Uncovered skills: {', '.join(uncovered)}")
        steps_log.append("   Consider increasing budget or adding specialized talents")
    else:
        steps_log.append("🎉 FULL COVERAGE ACHIEVED!")
        
    steps_log.append("═" * 50)
    steps_log.append(f"💵 Total Cost: ${total_cost:.3f} / ${budget:.3f}")
    steps_log.append(f"👥 Team Size: {len(recruited)} talents")
    
    result_json = {
        "task": task["description"],
        "recruited_talents": recruited,
        "total_cost": round(total_cost, 4),
        "budget": budget,
        "coverage_ratio": len(covered_skills & required_skills) / len(required_skills),
        "success": len(uncovered) == 0
    }
    
    return "\n".join(steps_log), json.dumps(result_json, indent=2)

def simulate_e2r_tree(task_description: str, exploration_width: int, max_depth: int) -> str:
    steps = [f"🌳 E²R TREE SEARCH SIMULATION",
             f"Task: {task_description}",
             f"Parameters: width={exploration_width}, max_depth={max_depth}",
             "═" * 50]
    
    steps.append("\n📥 PHASE 1: EXPLORE (Top-down decomposition)")
    
    subtasks = [
        {"name": f"Subtask_{i+1}", 
         "estimated_difficulty": random.choice(["low", "medium", "high"]),
         "candidates": min(exploration_width, 3 + i)}
        for i in range(min(5, max_depth * 2))
    ]
    
    for st in subtasks:
        steps.append(f"   └─ {st['name']} [{st['estimated_difficulty']}] → {st['candidates']} candidate approaches")
    
    steps.append("\n⚡ PHASE 2: EXECUTE (Execution & branching)")
    
    total_attempts = 0
    successful = 0
    
    for st in subtasks:
        attempts = min(st["candidates"], exploration_width)
        for j in range(attempts):
            total_attempts += 1
            success_prob = 0.7 if st["estimated_difficulty"] == "low" else (0.5 if st["estimated_difficulty"] == "medium" else 0.3)
            outcome = "✓" if random.random() < success_prob else "✗"
            if outcome == "✓":
                successful += 1
            steps.append(f"   ├─ {st['name']}/attempt_{j+1}: {outcome}")
    
    steps.append("\n🔍 PHASE 3: REVIEW (Bottom-up aggregation)")
    
    completion_rate = successful / total_attempts if total_attempts > 0 else 0
    
    if completion_rate >= 0.8:
        verdict = "TERMINATE_SUCCESS"
        steps.append(f"   └─ Aggregate success rate: {completion_rate:.1%}")
        steps.append(f"   └─ Verdict: {verdict}")
        steps.append(f"   └─ Output: Final deliverable compiled from {successful} successful subtask executions")
    elif completion_rate >= 0.5:
        verdict = "REFINE_PARTIAL"
        steps.append(f"   └─ Aggregate success rate: {completion_rate:.1%}")
        steps.append(f"   └─ Verdict: {verdict}")
        steps.append(f"   └─ Action: Retry failed {total_attempts - successful} subtasks with adjusted parameters")
    else:
        verdict = "REFINE_ALL"
        steps.append(f"   └─ Aggregate success rate: {completion_rate:.1%}")
        steps.append(f"   └─ Verdict: {verdict}")
        steps.append(f"   └─ Action: Backtrack to EXPLORE phase with wider width")
    
    steps.append("\n" + "═" * 50)
    steps.append(f"📊 Summary: {successful}/{total_attempts} attempts succeeded ({completion_rate:.1%})")
    steps.append(f"🔄 Termination guarantee: Tree depth bounded, deadlock-free by construction")
    
    return "\n".join(steps)

def get_talent_info() -> str:
    lines = ["📚 AVAILABLE TALENTS IN MARKET", "═" * 60]
    
    for tid, talent in TALENT_DATABASE.items():
        lines.append(f"\n🔹 {talent['name']} (ID: {tid})")
        lines.append(f"   Skills: {', '.join(talent['skills'])}")
        lines.append(f"   Tools: {', '.join(talent['tools'])}")
        lines.append(f"   Backend: {talent['backend']}")
        lines.append(f"   Cost: ${talent['cost_per_task']:.3f} | Reliability: {talent['reliability']:.0%}")
    
    return "\n".join(lines)

with gr.Blocks(title="OneManCompany Explorer") as demo:
    gr.Markdown("""
    # 🏢 OneManCompany (OMC) - Organizational Layer Demo
    
    Explore the framework from ["From Skills to Talent"](https://huggingface.co/papers/2604.22446) (Yu et al., 2026)
    
    This interactive demo illustrates two core OMC concepts:
    1. **Talent Market** - Dynamic recruitment of portable agent identities
    2. **E²R Tree Search** - Explore-Execute-Review hierarchical decision loop
    """)
    
    with gr.Tab("📋 Talent Market"):
        gr.Markdown("Simulate OMC's on-demand talent recruitment for capability gaps")
        
        task_dropdown = gr.Dropdown(
            choices=[(f"{t['id']}: {t['description'][:50]}...", i) for i, t in enumerate(TASK_TEMPLATES)],
            value=0,
            label="Select Task"
        )
        
        with gr.Row():
            with gr.Column(scale=2):
                recruit_btn = gr.Button("🎯 Recruit Optimal Team", variant="primary")
                talent_display = gr.Textbox(label="Available Talents", value=get_talent_info(), lines=15)
            
            with gr.Column(scale=3):
                recruitment_log = gr.Textbox(label="Recruitment Log", lines=12)
                json_output = gr.Textbox(label="Structured Result (JSON)", lines=8)
        
        recruit_btn.click(
            fn=recruit_talents_for_task,
            inputs=task_dropdown,
            outputs=[recruitment_log, json_output]
        )
    
    with gr.Tab("🌳 E²R Tree Search"):
        gr.Markdown("Simulate the Explore-Execute-Review hierarchical loop with termination guarantees")
        
        task_input = gr.Textbox(
            label="Task Description", 
            value="Build a web scraper that extracts product prices and alerts on changes",
            lines=2
        )
        
        with gr.Row():
            width_slider = gr.Slider(1, 5, value=3, step=1, label="Exploration Width (branching factor)")
            depth_slider = gr.Slider(1, 4, value=2, step=1, label="Max Tree Depth")
        
        run_e2r_btn = gr.Button("▶️ Run E²R Simulation", variant="primary")
        e2r_output = gr.Textbox(label="E²R Execution Trace", lines=25)
        
        run_e2r_btn.click(
            fn=simulate_e2r_tree,
            inputs=[task_input, width_slider, depth_slider],
            outputs=e2r_output
        )
    
    with gr.Tab("ℹ️ About OMC"):
        gr.Markdown("""
        ### Core Contributions from the Paper
        
        **OneManCompany** addresses a fundamental gap in multi-agent systems: the absence of a 
        principled *organizational layer* that governs how agent workforces are assembled, 
        governed, and improved over time.
        
        #### Key Innovations:
        
        1. **Talent Abstraction** - Encapsulates skills, tools, and runtime configs into 
           portable identities that abstract over heterogeneous backends
           
        2. **Talent Market** - Community-driven recruitment enabling on-demand capability 
           acquisition and dynamic team reconfiguration
           
        3. **E²R Tree Search** - Unified hierarchical loop combining:
           - **Explore**: Top-down task decomposition into accountable units
           - **Execute**: Mid-level plan execution with branching
           - **Review**: Bottom-up outcome aggregation and refinement
           
        4. **Formal Guarantees**: Termination and deadlock-freedom by construction
        
        #### Empirical Results:
        - **84.67%** success rate on PRDBench (surpassing SOTA by 15.48pp)
        - Cross-domain generalization demonstrated
        """)

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