"""
Advanced Team Management System
-----------------------------
Manages specialized teams of agents that work together towards common goals:
1. Team A: Coders (App/Software Developers)
2. Team B: Business (Entrepreneurs)
3. Team C: Research (Deep Online Research)
4. Team D: Crypto & Sports Trading

Features:
- Cross-team collaboration
- Goal alignment
- Resource sharing
- Synchronized execution
"""

from typing import Dict, List, Optional, Set, Union, TypeVar, Any
from dataclasses import dataclass, field
from enum import Enum
import asyncio
from datetime import datetime
import uuid
from collections import defaultdict

from orchestrator import AgentOrchestrator, TaskPriority, AgentRole, AgentState
from reasoning import UnifiedReasoningEngine

# Agent capabilities and personality types
class AgentCapability(Enum):
    """Core capabilities of agents."""
    REASONING = "reasoning"
    LEARNING = "learning"
    EXECUTION = "execution"
    COORDINATION = "coordination"
    MONITORING = "monitoring"

class AgentPersonality(Enum):
    """Different personality types for agents."""
    ANALYTICAL = "analytical"
    CREATIVE = "creative"
    PRAGMATIC = "pragmatic"
    COLLABORATIVE = "collaborative"
    PROACTIVE = "proactive"
    CAUTIOUS = "cautious"

class TeamType(Enum):
    """Specialized team types."""
    CODERS = "coders"
    BUSINESS = "business"
    RESEARCH = "research"
    TRADERS = "traders"

class TeamObjective(Enum):
    """Types of team objectives."""
    SOFTWARE_DEVELOPMENT = "software_development"
    BUSINESS_OPPORTUNITY = "business_opportunity"
    MARKET_RESEARCH = "market_research"
    TRADING_STRATEGY = "trading_strategy"
    CROSS_TEAM_PROJECT = "cross_team_project"

@dataclass
class TeamProfile:
    """Team profile and capabilities."""
    id: str
    type: TeamType
    name: str
    primary_objective: TeamObjective
    secondary_objectives: List[TeamObjective]
    agent_count: int
    expertise_areas: List[str]
    collaboration_score: float = 0.0
    success_rate: float = 0.0
    active_projects: int = 0

@dataclass
class CollaborationLink:
    """Defines collaboration between teams."""
    team_a_id: str
    team_b_id: str
    strength: float
    active_projects: int
    last_interaction: datetime
    success_rate: float

class TeamManager:
    """Manages specialized teams and their collaboration."""
    
    def __init__(self, orchestrator: AgentOrchestrator):
        self.orchestrator = orchestrator
        self.teams: Dict[str, TeamProfile] = {}
        self.agents: Dict[str, Dict[str, 'Agent']] = {}  # team_id -> {agent_id -> Agent}
        self.collaboration_network: Dict[str, CollaborationLink] = {}
        self.shared_objectives: Dict[str, Set[str]] = defaultdict(set)  # objective_id -> set of team_ids
        self.lock = asyncio.Lock()
        
        # Initialize specialized teams
        self._init_teams()

    def _init_teams(self):
        """Initialize specialized teams."""
        team_configs = {
            TeamType.CODERS: {
                "name": "Development Team",
                "primary": TeamObjective.SOFTWARE_DEVELOPMENT,
                "secondary": [
                    TeamObjective.BUSINESS_OPPORTUNITY,
                    TeamObjective.MARKET_RESEARCH
                ],
                "expertise": [
                    "full_stack_development",
                    "cloud_architecture",
                    "ai_ml",
                    "blockchain",
                    "mobile_development"
                ]
            },
            TeamType.BUSINESS: {
                "name": "Business Strategy Team",
                "primary": TeamObjective.BUSINESS_OPPORTUNITY,
                "secondary": [
                    TeamObjective.MARKET_RESEARCH,
                    TeamObjective.TRADING_STRATEGY
                ],
                "expertise": [
                    "market_analysis",
                    "business_strategy",
                    "digital_transformation",
                    "startup_innovation",
                    "product_management"
                ]
            },
            TeamType.RESEARCH: {
                "name": "Research & Analysis Team",
                "primary": TeamObjective.MARKET_RESEARCH,
                "secondary": [
                    TeamObjective.BUSINESS_OPPORTUNITY,
                    TeamObjective.TRADING_STRATEGY
                ],
                "expertise": [
                    "deep_research",
                    "data_analysis",
                    "trend_forecasting",
                    "competitive_analysis",
                    "technology_assessment"
                ]
            },
            TeamType.TRADERS: {
                "name": "Trading & Investment Team",
                "primary": TeamObjective.TRADING_STRATEGY,
                "secondary": [
                    TeamObjective.MARKET_RESEARCH,
                    TeamObjective.BUSINESS_OPPORTUNITY
                ],
                "expertise": [
                    "crypto_trading",
                    "sports_betting",
                    "risk_management",
                    "market_timing",
                    "portfolio_optimization"
                ]
            }
        }
        
        for team_type, config in team_configs.items():
            team_id = str(uuid.uuid4())
            self.teams[team_id] = TeamProfile(
                id=team_id,
                type=team_type,
                name=config["name"],
                primary_objective=config["primary"],
                secondary_objectives=config["secondary"],
                agent_count=5,  # Default size
                expertise_areas=config["expertise"]
            )
            self.agents[team_id] = {}

    async def initialize_team_agents(self):
        """Initialize agents for each team with appropriate roles and capabilities."""
        for team_id, team in self.teams.items():
            await self._create_team_agents(team_id)
            await self._establish_collaboration_links(team_id)

    async def _create_team_agents(self, team_id: str):
        """Create specialized agents for a team."""
        team = self.teams[team_id]
        
        # Define agent configurations based on team type
        agent_configs = self._get_agent_configs(team.type)
        
        for config in agent_configs:
            agent_id = await self.orchestrator.create_agent(
                role=config["role"],
                capabilities=config["capabilities"]
            )
            
            agent = Agent(
                profile=config["profile"],
                reasoning_engine=self.orchestrator.reasoning_engine,
                meta_learning=self.orchestrator.meta_learning,
                config=config.get("config", {})
            )
            
            self.agents[team_id][agent_id] = agent

    def _get_agent_configs(self, team_type: TeamType) -> List[Dict]:
        """Get agent configurations based on team type."""
        base_configs = [
            {
                "role": AgentRole.COORDINATOR,
                "capabilities": [
                    AgentCapability.REASONING,
                    AgentCapability.COORDINATION
                ],
                "personality": AgentPersonality.PROACTIVE,
                "profile": {
                    "name": "Coordinator",
                    "description": "Team coordinator"
                }
            },
            {
                "role": AgentRole.EXECUTOR,
                "capabilities": [
                    AgentCapability.EXECUTION,
                    AgentCapability.LEARNING
                ],
                "personality": AgentPersonality.ANALYTICAL,
                "profile": {
                    "name": "Executor",
                    "description": "Task executor"
                }
            }
        ]
        
        # Add team-specific configurations
        if team_type == TeamType.CODERS:
            base_configs.extend([
                {
                    "role": AgentRole.EXECUTOR,
                    "capabilities": [
                        AgentCapability.EXECUTION,
                        AgentCapability.REASONING
                    ],
                    "personality": AgentPersonality.CREATIVE,
                    "expertise": ["software_development", "system_design"],
                    "profile": {
                        "name": "Developer",
                        "description": "Software developer"
                    }
                }
            ])
        elif team_type == TeamType.BUSINESS:
            base_configs.extend([
                {
                    "role": AgentRole.PLANNER,
                    "capabilities": [
                        AgentCapability.REASONING,
                        AgentCapability.LEARNING
                    ],
                    "personality": AgentPersonality.PROACTIVE,
                    "expertise": ["business_strategy", "market_analysis"],
                    "profile": {
                        "name": "Planner",
                        "description": "Business planner"
                    }
                }
            ])
        elif team_type == TeamType.RESEARCH:
            base_configs.extend([
                {
                    "role": AgentRole.MONITOR,
                    "capabilities": [
                        AgentCapability.MONITORING,
                        AgentCapability.LEARNING
                    ],
                    "personality": AgentPersonality.ANALYTICAL,
                    "expertise": ["research", "data_analysis"],
                    "profile": {
                        "name": "Researcher",
                        "description": "Researcher"
                    }
                }
            ])
        elif team_type == TeamType.TRADERS:
            base_configs.extend([
                {
                    "role": AgentRole.EXECUTOR,
                    "capabilities": [
                        AgentCapability.EXECUTION,
                        AgentCapability.REASONING
                    ],
                    "personality": AgentPersonality.CAUTIOUS,
                    "expertise": ["trading", "risk_management"],
                    "profile": {
                        "name": "Trader",
                        "description": "Trader"
                    }
                }
            ])
            
        return base_configs

    async def _establish_collaboration_links(self, team_id: str):
        """Establish collaboration links with other teams."""
        team = self.teams[team_id]
        
        for other_id, other_team in self.teams.items():
            if other_id != team_id:
                link_id = f"{min(team_id, other_id)}_{max(team_id, other_id)}"
                if link_id not in self.collaboration_network:
                    self.collaboration_network[link_id] = CollaborationLink(
                        team_a_id=team_id,
                        team_b_id=other_id,
                        strength=0.5,  # Initial collaboration strength
                        active_projects=0,
                        last_interaction=datetime.now(),
                        success_rate=0.0
                    )

    async def create_cross_team_objective(
        self,
        objective: str,
        required_teams: List[TeamType],
        priority: TaskPriority = TaskPriority.MEDIUM
    ) -> str:
        """Create an objective that requires multiple teams."""
        objective_id = str(uuid.uuid4())
        
        # Find relevant teams
        selected_teams = []
        for team_id, team in self.teams.items():
            if team.type in required_teams:
                selected_teams.append(team_id)
        
        if len(selected_teams) < len(required_teams):
            raise ValueError("Not all required teams are available")
        
        # Create shared objective
        self.shared_objectives[objective_id].update(selected_teams)
        
        # Create tasks for each team
        tasks = []
        for team_id in selected_teams:
            task_id = await self.orchestrator.submit_task(
                description=f"Team {self.teams[team_id].name} contribution to: {objective}",
                priority=priority
            )
            tasks.append(task_id)
        
        return objective_id

    async def monitor_objective_progress(self, objective_id: str) -> Dict:
        """Monitor progress of a cross-team objective."""
        if objective_id not in self.shared_objectives:
            raise ValueError("Unknown objective")
            
        team_progress = {}
        for team_id in self.shared_objectives[objective_id]:
            team = self.teams[team_id]
            team_agents = self.agents[team_id]
            
            # Calculate team progress
            active_agents = sum(1 for agent in team_agents.values() if agent.state == AgentState.BUSY)
            completion_rate = sum(agent.get_task_completion_rate() for agent in team_agents.values()) / len(team_agents)
            
            team_progress[team.name] = {
                "active_agents": active_agents,
                "completion_rate": completion_rate,
                "collaboration_score": team.collaboration_score
            }
            
        return team_progress

    async def optimize_team_collaboration(self):
        """Optimize collaboration between teams."""
        for link in self.collaboration_network.values():
            team_a = self.teams[link.team_a_id]
            team_b = self.teams[link.team_b_id]
            
            # Update collaboration strength based on:
            # 1. Number of successful joint projects
            # 2. Frequency of interaction
            # 3. Complementary expertise
            
            success_factor = link.success_rate
            interaction_factor = min((datetime.now() - link.last_interaction).days / 30.0, 1.0)
            expertise_overlap = len(
                set(team_a.expertise_areas) & set(team_b.expertise_areas)
            ) / len(set(team_a.expertise_areas) | set(team_b.expertise_areas))
            
            new_strength = (
                0.4 * success_factor +
                0.3 * (1 - interaction_factor) +
                0.3 * (1 - expertise_overlap)
            )
            
            link.strength = 0.7 * link.strength + 0.3 * new_strength

    async def get_team_recommendations(self, objective: str) -> List[TeamType]:
        """Get recommended teams for an objective based on expertise and collaboration history."""
        # Analyze objective to determine required expertise
        required_expertise = await self._analyze_objective(objective)
        
        # Score each team
        team_scores = {}
        for team_id, team in self.teams.items():
            # Calculate expertise match
            expertise_match = len(
                set(required_expertise) & set(team.expertise_areas)
            ) / len(required_expertise)
            
            # Calculate collaboration potential
            collab_potential = self._calculate_collaboration_potential(team_id)
            
            # Calculate success history
            success_history = team.success_rate
            
            # Weighted score
            score = (
                0.4 * expertise_match +
                0.3 * collab_potential +
                0.3 * success_history
            )
            
            team_scores[team.type] = score
            
        # Return sorted recommendations
        return sorted(
            team_scores.keys(),
            key=lambda x: team_scores[x],
            reverse=True
        )

    async def _analyze_objective(self, objective: str) -> List[str]:
        """Analyze an objective to determine required expertise."""
        # Use reasoning engine to analyze objective
        analysis = await self.orchestrator.reasoning_engine.reason(
            query=f"Analyze required expertise for: {objective}",
            context={
                "available_expertise": [
                    expertise
                    for team in self.teams.values()
                    for expertise in team.expertise_areas
                ]
            }
        )
        
        return analysis.get("required_expertise", [])

    def _calculate_collaboration_potential(self, team_id: str) -> float:
        """Calculate a team's collaboration potential based on history."""
        team_links = [
            link for link in self.collaboration_network.values()
            if team_id in (link.team_a_id, link.team_b_id)
        ]
        
        if not team_links:
            return 0.5
            
        return sum(link.strength for link in team_links) / len(team_links)

    async def update_team_metrics(self):
        """Update performance metrics for all teams."""
        for team_id, team in self.teams.items():
            team_agents = self.agents[team_id]
            
            # Calculate success rate
            completed_tasks = sum(
                agent.get_completed_task_count()
                for agent in team_agents.values()
            )
            total_tasks = sum(
                agent.get_total_task_count()
                for agent in team_agents.values()
            )
            team.success_rate = completed_tasks / max(1, total_tasks)
            
            # Calculate collaboration score
            team_links = [
                link for link in self.collaboration_network.values()
                if team_id in (link.team_a_id, link.team_b_id)
            ]
            team.collaboration_score = (
                sum(link.strength for link in team_links) /
                len(team_links) if team_links else 0.5
            )

class Agent:
    def __init__(self, profile: Dict, reasoning_engine: UnifiedReasoningEngine, meta_learning: bool, config: Optional[Dict[str, Any]] = None):
        self.profile = profile
        self.config = config or {}
        
        # Use provided reasoning engine or create one with config
        self.reasoning_engine = reasoning_engine if reasoning_engine else UnifiedReasoningEngine(
            min_confidence=self.config.get('min_confidence', 0.7),
            parallel_threshold=self.config.get('parallel_threshold', 3),
            learning_rate=self.config.get('learning_rate', 0.1),
            strategy_weights=self.config.get('strategy_weights', {
                "LOCAL_LLM": 0.8,
                "CHAIN_OF_THOUGHT": 0.6,
                "TREE_OF_THOUGHTS": 0.5,
                "META_LEARNING": 0.4
            })
        )
        self.meta_learning = meta_learning
        self.state = AgentState.IDLE

    def get_task_completion_rate(self):
        # Implement task completion rate calculation
        pass

    def get_completed_task_count(self):
        # Implement completed task count calculation
        pass

    def get_total_task_count(self):
        # Implement total task count calculation
        pass