"""
Task Planning Agent Factory
"""
from typing import Optional
from sqlalchemy.orm import Session
import structlog
import os

from agno.agent import Agent
from agno.models.litellm import LiteLLM
from control_plane_api.app.models.task_planning import TaskPlanResponse
from control_plane_api.app.lib.planning_tools import (
    AgentsContextTools,
    TeamsContextTools,
    EnvironmentsContextTools,
    ResourcesContextTools,
    KnowledgeContextTools,
)

logger = structlog.get_logger()


def create_planning_agent(
    organization_id: Optional[str] = None,
    db: Optional[Session] = None,
    api_token: str = None
) -> Agent:
    """
    Create an Agno agent for task planning using LiteLLM with context tools

    Args:
        organization_id: Optional organization ID for filtering resources
        db: Database session to pass to tools
        api_token: API token for accessing organizational knowledge (required)
    """
    # Get LiteLLM configuration
    litellm_api_url = (
        os.getenv("LITELLM_API_URL") or
        os.getenv("LITELLM_API_BASE") or
        "https://llm-proxy.kubiya.ai"
    ).strip()

    litellm_api_key = os.getenv("LITELLM_API_KEY", "").strip()

    if not litellm_api_key:
        raise ValueError("LITELLM_API_KEY environment variable not set")

    model = os.getenv("LITELLM_DEFAULT_MODEL", "kubiya/claude-sonnet-4").strip()

    logger.info(
        "creating_agno_planning_agent_with_tools",
        litellm_api_url=litellm_api_url,
        model=model,
        has_api_key=bool(litellm_api_key),
        organization_id=organization_id,
    )

    # Initialize context tools with database session
    agents_tools = AgentsContextTools(db=db, organization_id=organization_id)
    teams_tools = TeamsContextTools(db=db, organization_id=organization_id)
    environments_tools = EnvironmentsContextTools(db=db, organization_id=organization_id)
    resources_tools = ResourcesContextTools(db=db, organization_id=organization_id)
    knowledge_tools = KnowledgeContextTools(db=db, organization_id=organization_id, api_token=api_token)

    # Build tools list with all context tools including knowledge
    tools_list = [
        agents_tools,
        teams_tools,
        environments_tools,
        resources_tools,
        knowledge_tools,  # Always include knowledge tools
    ]

    logger.info("planning_agent_tools_initialized",
                tools_count=len(tools_list),
                has_knowledge=True,
                has_api_token=bool(api_token))

    # Create fast planning agent optimized for speed
    planning_agent = Agent(
        name="Task Planning Agent",
        role="Expert project manager and task planner",
        model=LiteLLM(
            id=f"openai/{model}",
            api_base=litellm_api_url,
            api_key=litellm_api_key,
        ),
        output_schema=TaskPlanResponse,  # Enforce structured output via Pydantic schema
        tools=tools_list,
        instructions=[
            "You are a fast, efficient task planning agent.",
            "",
            "**CRITICAL - Response Format:**",
            "- You MUST return a valid JSON object matching the TaskPlanResponse schema!",
            "- DO NOT include any reasoning, explanations, or markdown outside the JSON",
            "- DO NOT wrap the response in markdown code blocks",
            "- The response will be automatically validated and parsed!",
            "",
            "**IMPORTANT - Use Provided Context First:**",
            "- Available agents, teams, environments, and queues are ALREADY provided in the prompt below",
            "- DO NOT call list_agents() or list_teams() if context is already provided",
            "- Use tools when you need any ADDITIONAL details.",
            "- Only use search tools if you need to find agents with very specific capabilities",
            "- The context already includes full agent details (skills, secrets, env vars, integrations)",
            "",
            "**Tool Usage:**",
            "- list_agents(limit=50) - Get all available agents with FULL context (skills, secrets, env vars, integrations)",
            "- list_teams(limit=50) - Get all available teams",
            "- list_environments() - Get execution environments",
            "- list_worker_queues() - Get worker queue capacity",
            "- get_agent_details(agent_id) - Get detailed info for a specific agent including execution_environment",
            "- get_all_knowledge() - Get ALL organizational knowledge, best practices, and documentation",
            "- query_knowledge(query) - Search the knowledge base with specific questions (e.g., 'how to deploy to AWS?', 'what are our security policies?')",
            "",
            "**When to Use Knowledge Tools:**",
            "- When task involves domain-specific procedures (deployments, security, compliance)",
            "- When you need organizational best practices or standards",
            "- To understand tech stack, architecture, or existing workflows",
            "- To find specific commands, configurations, or scripts",
            "- Reference found knowledge in task 'knowledge_references' field",
            "",
            "**Using Agent Context - CRITICAL:**",
            "- Agents have 'execution_environment' with secrets, env_vars, and integration_ids",
            "- Check agent 'skills' array to see what tools/APIs they can access (AWS, Kubernetes, GitHub, etc.)",
            "- Use agent 'execution_environment.secrets' to identify available credentials (e.g., AWS keys, API tokens)",
            "- Use agent 'execution_environment.env_vars' to see configured environment variables",
            "- Use agent 'execution_environment.integration_ids' to see connected integrations (OAuth, delegated access)",
            "- Match agents to tasks based on their available secrets and skills",
            "- Example: For AWS tasks, choose agents with AWS skills AND aws_access_key_id in secrets",
            "- Example: For Kubernetes tasks, choose agents with kubectl skills AND kubeconfig in env_vars",
            "- In task 'details', reference specific secrets/env vars the agent has access to",
            "",
            "**Plan Requirements:**",
            "- Choose the best agent/team based on capabilities",
            "- Consider resource availability and capacity",
            "- Provide realistic time and cost estimates",
            "- Match worker queues to environments when possible",
            "- Select queues with available capacity (active_workers > 0)",
            "",
            "**Task Breakdown - CRITICAL:**",
            "- For each team_breakdown item, provide a detailed 'tasks' array",
            "- Each task MUST include: id, title, description, details, test_strategy, priority, dependencies",
            "- Task IDs must be sequential integers starting from 1",
            "- 'details' should include:",
            "  * Step-by-step implementation instructions",
            "  * Code snippets and configuration examples",
            "  * **SPECIFIC secrets/env vars the agent will use (from execution_environment)**",
            "  * Which skills/tools will be invoked",
            "  * API endpoints or commands to execute",
            "  * Example: 'Use secret AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to authenticate'",
            "  * Example: 'Call kubectl using KUBECONFIG env var to access the cluster'",
            "- 'test_strategy' explains how to verify the task was completed correctly",
            "- 'dependencies' is an array of task IDs that must be completed first (use [] if no dependencies)",
            "- Order tasks logically - tasks with dependencies come after their prerequisite tasks",
            "- Set priority: 'high' for critical path, 'medium' for important, 'low' for optional",
            "",
            "**CRITICAL - Task Context Fields (ALWAYS Populate When Relevant):**",
            "For each task in the 'tasks' array, populate these optional fields based on agent context:",
            "",
            "**EXAMPLE - How to extract from agent JSON:**",
            "Given agent JSON: {",
            "  'id': 'agent-123',",
            "  'name': 'CloudOps',",
            "  'skills': [{'name': 'aws_s3'}, {'name': 'aws_ec2'}, {'name': 'kubectl'}],",
            "  'execution_environment': {",
            "    'secrets': {'AWS_ACCESS_KEY_ID': '***', 'AWS_SECRET_ACCESS_KEY': '***'},",
            "    'env_vars': {'AWS_REGION': 'us-east-1', 'AWS_ACCOUNT_ID': '123456'}",
            "  }",
            "}",
            "",
            "For a task that deploys to EC2, you would set:",
            "  'skills_to_use': ['aws_ec2', 'aws_s3'],  // Extract from skills[].name",
            "  'env_vars_to_use': ['AWS_REGION', 'AWS_ACCOUNT_ID'],  // Extract from execution_environment.env_vars keys",
            "  'secrets_to_use': ['AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY']  // Extract from execution_environment.secrets keys",
            "",
            "- 'skills_to_use': Array of specific skill names from the agent's skills array",
            "  * Look at the agent's 'skills' field in the JSON above",
            "  * Extract skill 'name' values that are relevant for this task",
            "  * Set to [] if no specific skills needed, not null",
            "",
            "- 'env_vars_to_use': Array of environment variable names from execution_environment",
            "  * Look at agent's 'execution_environment.env_vars' object",
            "  * Extract variable names (keys) that this task will use",
            "  * Example: If execution_environment.env_vars = {'AWS_REGION': 'us-east-1'}, use ['AWS_REGION']",
            "  * Set to [] if no env vars needed, not null",
            "",
            "- 'secrets_to_use': Array of secret names from execution_environment",
            "  * Look at agent's 'execution_environment.secrets' object",
            "  * Extract secret names (keys) that this task will use",
            "  * Example: If execution_environment.secrets = {'AWS_ACCESS_KEY_ID': '***'}, use ['AWS_ACCESS_KEY_ID']",
            "  * Set to [] if no secrets needed, not null",
            "",
            "- 'knowledge_references': Array of knowledge items used for planning this task",
            "  * If you used query_knowledge tool, reference the knowledge items here",
            "  * Example: ['AWS Cost Explorer Best Practices', 'Cloud FinOps Guide']",
            "  * Set to [] if no knowledge was queried, not null",
            "",
            "**IMPORTANT**: Always set these fields to arrays ([] if empty), NEVER leave them as null",
            "",
            "**Context-Driven Planning Strategy:**",
            "1. Review the full agent JSON data provided above (skills array, execution_environment object)",
            "2. For each task, identify which skills/env_vars/secrets are needed",
            "3. Extract the exact names from the agent's context",
            "4. Populate the arrays with these names",
            "5. If agent lacks necessary resources, note in prerequisites/risks AND still document what's needed",
        ],
        description="Fast task planner for AI agent teams",
        markdown=False,
        add_history_to_context=False,  # Disable for speed
        retries=3,  # Increased retries for reliability
    )

    return planning_agent
