"""Team-related Temporal activities"""

import os
from dataclasses import dataclass
from typing import Optional, List
from temporalio import activity
import structlog

from agno.agent import Agent
from agno.team import Team
from agno.models.litellm import LiteLLM

from control_plane_api.app.lib.supabase import get_supabase
from control_plane_api.app.activities.agent_activities import update_execution_status, ActivityUpdateExecutionInput

logger = structlog.get_logger()


@dataclass
class ActivityGetTeamAgentsInput:
    """Input for get_team_agents activity"""
    team_id: str
    organization_id: str


@dataclass
class ActivityExecuteTeamInput:
    """Input for execute_team_coordination activity"""
    execution_id: str
    team_id: str
    organization_id: str
    prompt: str
    system_prompt: Optional[str] = None
    agents: List[dict] = None
    team_config: dict = None

    def __post_init__(self):
        if self.agents is None:
            self.agents = []
        if self.team_config is None:
            self.team_config = {}


@activity.defn
async def get_team_agents(input: ActivityGetTeamAgentsInput) -> dict:
    """
    Get all agents in a team.

    This activity fetches agents belonging to the team from the database.

    Args:
        input: Activity input with team details

    Returns:
        Dict with agents list
    """
    print(f"\n\n=== GET_TEAM_AGENTS START ===")
    print(f"team_id: {input.team_id} (type: {type(input.team_id).__name__})")
    print(f"organization_id: {input.organization_id} (type: {type(input.organization_id).__name__})")
    print(f"================================\n")

    activity.logger.info(
        f"[DEBUG] Getting team agents START",
        extra={
            "team_id": input.team_id,
            "team_id_type": type(input.team_id).__name__,
            "organization_id": input.organization_id,
            "organization_id_type": type(input.organization_id).__name__,
        }
    )

    try:
        client = get_supabase()
        print(f"Supabase client created successfully")

        activity.logger.info(
            f"[DEBUG] Supabase client created, fetching team configuration",
            extra={
                "team_id": input.team_id,
                "organization_id": input.organization_id,
            }
        )

        # First, get the team configuration to get member_ids (source of truth)
        print(f"Fetching team configuration for team_id={input.team_id}")
        team_result = (
            client.table("teams")
            .select("configuration")
            .eq("id", input.team_id)
            .eq("organization_id", input.organization_id)
            .single()
            .execute()
        )

        if not team_result.data:
            print(f"Team not found!")
            activity.logger.error(
                f"[DEBUG] Team not found",
                extra={
                    "team_id": input.team_id,
                    "organization_id": input.organization_id,
                }
            )
            return {"agents": [], "count": 0}

        team_config = team_result.data.get("configuration", {})
        member_ids = team_config.get("member_ids", [])
        print(f"Team configuration member_ids: {member_ids}")

        activity.logger.info(
            f"[DEBUG] Team configuration loaded",
            extra={
                "team_id": input.team_id,
                "member_ids": member_ids,
                "member_count": len(member_ids),
            }
        )

        # Get agents by member_ids (not by FK relationship)
        agents = []
        if member_ids:
            print(f"Fetching agents with IDs: {member_ids}")
            result = (
                client.table("agents")
                .select("*")
                .in_("id", member_ids)
                .eq("organization_id", input.organization_id)
                .execute()
            )
            agents = result.data or []
        print(f"Query executed. Result data length: {len(agents)}")

        activity.logger.info(
            f"[DEBUG] Query executed, processing results",
            extra={
                "agents_found": len(agents),
                "agent_ids": [a.get("id") for a in agents],
            }
        )

        print(f"Agents found: {len(agents)}")
        if agents:
            for agent in agents:
                print(f"  - {agent.get('name')} (ID: {agent.get('id')})")

        activity.logger.info(
            f"[DEBUG] Retrieved team agents",
            extra={
                "team_id": input.team_id,
                "agent_count": len(agents),
                "agent_names": [a.get("name") for a in agents],
                "agent_ids": [a.get("id") for a in agents],
            }
        )

        if not agents:
            print(f"\n!!! NO AGENTS FOUND - Running verification query !!!")
            activity.logger.warning(
                f"[DEBUG] WARNING: No agents found for team - running verification query",
                extra={
                    "team_id": input.team_id,
                    "organization_id": input.organization_id,
                }
            )

            # Try query without org filter to debug
            result_no_org = (
                client.table("agents")
                .select("id, name, team_id, organization_id")
                .eq("team_id", input.team_id)
                .execute()
            )
            print(f"Query without org filter returned {len(result_no_org.data or [])} agents")
            if result_no_org.data:
                for agent in result_no_org.data:
                    print(f"  - {agent.get('name')} (team_id: {agent.get('team_id')}, org_id: {agent.get('organization_id')})")

            activity.logger.warning(
                f"[DEBUG] Query without org filter returned {len(result_no_org.data or [])} agents",
                extra={
                    "agents_found": result_no_org.data if result_no_org.data else [],
                }
            )

        print(f"\n=== GET_TEAM_AGENTS END: Returning {len(agents)} agents ===\n\n")
        return {
            "agents": agents,
            "count": len(agents),
        }

    except Exception as e:
        print(f"\n!!! EXCEPTION in get_team_agents: {type(e).__name__}: {str(e)} !!!\n")
        activity.logger.error(
            f"[DEBUG] EXCEPTION in get_team_agents",
            extra={
                "team_id": input.team_id,
                "organization_id": input.organization_id,
                "error": str(e),
                "error_type": type(e).__name__,
            }
        )
        raise


@activity.defn
async def execute_team_coordination(input: ActivityExecuteTeamInput) -> dict:
    """
    Execute team coordination using Agno Teams.

    This activity creates an Agno Team with member Agents and executes
    the team run, allowing Agno to handle coordination.

    Args:
        input: Activity input with team execution details

    Returns:
        Dict with aggregated response, usage, success flag
    """
    activity.logger.info(
        f"Executing team coordination with Agno Teams",
        extra={
            "execution_id": input.execution_id,
            "team_id": input.team_id,
            "agent_count": len(input.agents),
        }
    )

    try:
        # Create Agno Agent objects for each team member
        member_agents = []
        for agent_data in input.agents:
            # Get model ID (default to kubiya/claude-sonnet-4 if not specified)
            model_id = agent_data.get("model_id") or "kubiya/claude-sonnet-4"

            # Get LiteLLM configuration from environment
            litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
            litellm_api_key = os.getenv("LITELLM_API_KEY")

            # Create Agno Agent with explicit LiteLLM proxy configuration
            # IMPORTANT: Use openai/ prefix for custom proxy compatibility (same as agno_service)
            member_agent = Agent(
                name=agent_data["name"],
                role=agent_data.get("description", agent_data["name"]),
                model=LiteLLM(
                    id=f"openai/{model_id}",  # e.g., "openai/kubiya/claude-sonnet-4"
                    api_base=litellm_api_base,
                    api_key=litellm_api_key,
                ),
            )
            member_agents.append(member_agent)

            activity.logger.info(
                f"Created Agno Agent",
                extra={
                    "agent_name": agent_data["name"],
                    "model": model_id,
                }
            )

        # Create Agno Team with member agents and LiteLLM model for coordination
        litellm_api_base = os.getenv("LITELLM_API_BASE", "https://llm-proxy.kubiya.ai")
        litellm_api_key = os.getenv("LITELLM_API_KEY")

        # Get coordinator model from team configuration (if specified by user in UI)
        # Falls back to default if not configured
        team_model = (
            input.team_config.get("llm", {}).get("model")
            or "kubiya/claude-sonnet-4"  # Default coordinator model
        )

        # Create Team with openai/ prefix for custom proxy compatibility (same as agno_service)
        team = Team(
            members=member_agents,
            name=f"Team {input.team_id}",
            model=LiteLLM(
                id=f"openai/{team_model}",  # e.g., "openai/kubiya/claude-sonnet-4"
                api_base=litellm_api_base,
                api_key=litellm_api_key,
            ),
        )

        activity.logger.info(
            f"Created Agno Team with {len(member_agents)} members",
            extra={
                "coordinator_model": team_model,
                "member_count": len(member_agents),
            }
        )

        # Execute team run in a thread pool since team.run() is synchronous
        # This prevents blocking the async event loop in Temporal
        import asyncio
        result = await asyncio.to_thread(team.run, input.prompt)

        activity.logger.info(
            f"Agno Team execution completed",
            extra={
                "execution_id": input.execution_id,
                "has_content": bool(result.content),
            }
        )

        # Extract response content
        response_content = result.content if hasattr(result, "content") else str(result)

        # Extract usage metrics if available
        usage = {}
        if hasattr(result, "metrics") and result.metrics:
            metrics = result.metrics
            usage = {
                "input_tokens": getattr(metrics, "input_tokens", 0),
                "output_tokens": getattr(metrics, "output_tokens", 0),
                "total_tokens": getattr(metrics, "total_tokens", 0),
            }

        return {
            "success": True,
            "response": response_content,
            "usage": usage,
            "coordination_type": "agno_team",
        }

    except Exception as e:
        activity.logger.error(
            f"Team coordination failed",
            extra={
                "execution_id": input.execution_id,
                "error": str(e),
            }
        )
        return {
            "success": False,
            "error": str(e),
            "coordination_type": "agno_team",
            "usage": {},
        }


async def _execute_sequential(input: ActivityExecuteTeamInput) -> dict:
    """Execute agents sequentially"""
    activity.logger.info("Executing team sequentially")

    responses = []
    total_usage = {}

    for agent in input.agents:
        try:
            # Import here to avoid circular dependency
            from control_plane_api.app.services.litellm_service import litellm_service

            model = agent.get("model_id") or "kubiya/claude-sonnet-4"
            model_config = agent.get("model_config", {})

            # Execute agent
            result = litellm_service.execute_agent(
                prompt=f"{input.prompt}\n\nAgent role: {agent.get('description', agent['name'])}",
                model=model,
                system_prompt=input.system_prompt,
                **model_config
            )

            if result.get("success"):
                responses.append({
                    "agent": agent["name"],
                    "response": result.get("response"),
                })

            # Aggregate usage
            usage = result.get("usage", {})
            for key, value in usage.items():
                total_usage[key] = total_usage.get(key, 0) + value

        except Exception as e:
            activity.logger.error(
                f"Agent execution failed",
                extra={"agent_id": agent["id"], "error": str(e)}
            )
            responses.append({
                "agent": agent["name"],
                "error": str(e),
            })

    # Combine responses
    combined_response = "\n\n".join([
        f"**{r['agent']}**: {r.get('response', r.get('error'))}"
        for r in responses
    ])

    return {
        "success": True,
        "response": combined_response,
        "usage": total_usage,
        "coordination_type": "sequential",
    }


async def _execute_parallel(input: ActivityExecuteTeamInput) -> dict:
    """Execute agents in parallel (placeholder - would use asyncio.gather in real impl)"""
    activity.logger.info("Executing team in parallel")
    # For now, fall back to sequential
    # In a real implementation, this would use asyncio.gather or child workflows
    return await _execute_sequential(input)


async def _execute_hierarchical(input: ActivityExecuteTeamInput) -> dict:
    """Execute agents hierarchically with a coordinator (placeholder)"""
    activity.logger.info("Executing team hierarchically")
    # For now, fall back to sequential
    # In a real implementation, this would have a coordinator agent that delegates
    return await _execute_sequential(input)
