"""PraisonAI Task Component for Langflow.

Creates a PraisonAI Task with structured output, workflow branching,
and conditional routing support.
"""

from __future__ import annotations

from typing import Any

from langflow.custom import Component
from langflow.io import (
    BoolInput,
    DictInput,
    DropdownInput,
    FileInput,
    HandleInput,
    IntInput,
    MessageTextInput,
    MultilineInput,
    Output,
)


class PraisonAITaskComponent(Component):
    """PraisonAI Task component for Langflow workflows.

    Creates a task definition that can be assigned to agents within
    a multi-agent orchestration workflow.
    """

    display_name: str = "PraisonAI Task"
    description: str = "Define a task for PraisonAI agent workflows with structured output and routing."
    documentation: str = "https://docs.praison.ai"
    icon: str = "bot"
    name: str = "PraisonAITask"

    inputs = [
        # ============================================================
        # CORE TASK DEFINITION
        # ============================================================
        MessageTextInput(
            name="task_name",
            display_name="Task Name",
            info="Unique name for this task.",
            value="task",
        ),
        MultilineInput(
            name="description",
            display_name="Description",
            info="Detailed description of what the task should accomplish.",
            value="Complete the assigned task.",
        ),
        MultilineInput(
            name="expected_output",
            display_name="Expected Output",
            info="Description of the expected output format.",
            value="A detailed response.",
        ),
        # ============================================================
        # AGENT ASSIGNMENT
        # ============================================================
        HandleInput(
            name="agent",
            display_name="Agent",
            info="Agent assigned to execute this task.",
            input_types=["Agent"],
        ),
        # ============================================================
        # CONTEXT (Task Dependencies)
        # ============================================================
        HandleInput(
            name="context",
            display_name="Context Tasks",
            info="Previous tasks whose output should be available as context.",
            input_types=["Task"],
            is_list=True,
            advanced=True,
        ),
        BoolInput(
            name="retain_full_context",
            display_name="Retain Full Context",
            info="Keep full context from previous tasks (increases token usage).",
            value=False,
            advanced=True,
        ),
        # ============================================================
        # STRUCTURED OUTPUT
        # ============================================================
        MultilineInput(
            name="output_json",
            display_name="Output JSON Schema",
            info='JSON schema for structured output (e.g., {"summary": "str", "score": "int"}).',
            advanced=True,
        ),
        MessageTextInput(
            name="output_file",
            display_name="Output File",
            info="File path to save task output.",
            advanced=True,
        ),
        BoolInput(
            name="create_directory",
            display_name="Create Directory",
            info="Create parent directories for output file if they don't exist.",
            value=True,
            advanced=True,
        ),
        # ============================================================
        # INPUT FILES & IMAGES
        # ============================================================
        FileInput(
            name="input_file",
            display_name="Input File",
            info="File to use as input for this task.",
            advanced=True,
        ),
        HandleInput(
            name="images",
            display_name="Images",
            info="Images to include with the task (for vision models).",
            input_types=["str"],
            is_list=True,
            advanced=True,
        ),
        # ============================================================
        # TOOLS (Task-specific)
        # ============================================================
        HandleInput(
            name="tools",
            display_name="Tools",
            info="Task-specific tools (supplements agent tools).",
            input_types=["Tool", "BaseTool"],
            is_list=True,
            advanced=True,
        ),
        # ============================================================
        # GUARDRAILS & VALIDATION
        # ============================================================
        BoolInput(
            name="guardrail",
            display_name="Guardrail",
            info="Enable output validation for this task.",
            value=False,
            advanced=True,
        ),
        IntInput(
            name="max_retries",
            display_name="Max Retries",
            info="Maximum retry attempts if task fails or validation fails.",
            value=3,
            advanced=True,
        ),
        BoolInput(
            name="quality_check",
            display_name="Quality Check",
            info="Enable quality metrics collection for task output.",
            value=True,
            advanced=True,
        ),
        # ============================================================
        # EXECUTION OPTIONS
        # ============================================================
        BoolInput(
            name="async_execution",
            display_name="Async Execution",
            info="Execute this task asynchronously.",
            value=False,
            advanced=True,
        ),
        BoolInput(
            name="rerun",
            display_name="Rerun",
            info="Allow task to be re-executed.",
            value=False,
            advanced=True,
        ),
        # ============================================================
        # WORKFLOW (Decision/Loop)
        # ============================================================
        DropdownInput(
            name="task_type",
            display_name="Task Type",
            info="Type of task: regular task, decision point, or loop.",
            options=["task", "decision", "loop"],
            value="task",
            advanced=True,
        ),
        DictInput(
            name="condition",
            display_name="Condition",
            info="Branching conditions for decision tasks (e.g., {'approved': ['task_a'], 'rejected': ['task_b']}).",
            advanced=True,
        ),
        MultilineInput(
            name="next_tasks",
            display_name="Next Tasks",
            info="Task names to execute after this task (comma-separated for workflow routing).",
            advanced=True,
        ),
        BoolInput(
            name="is_start",
            display_name="Is Start Task",
            info="Mark this as the starting task in a workflow.",
            value=False,
            advanced=True,
        ),
        # ============================================================
        # VARIABLES
        # ============================================================
        DictInput(
            name="variables",
            display_name="Variables",
            info="Variables for substitution in description (e.g., {'topic': 'AI'} replaces {topic}).",
            advanced=True,
        ),
    ]

    outputs = [
        Output(
            display_name="Task",
            name="task",
            method="build_task",
        ),
    ]

    def _import_task(self):
        """Import Task class with proper error handling."""
        try:
            from praisonaiagents import Task
        except ImportError as e:
            msg = "PraisonAI Agents is not installed. Install with: pip install praisonaiagents"
            raise ImportError(msg) from e
        else:
            return Task

    def _parse_output_json(self):
        """Parse JSON schema string into a Pydantic model if provided."""
        if not self.output_json:
            return None

        try:
            import json

            from pydantic import create_model

            schema = json.loads(self.output_json)

            # Create a dynamic Pydantic model from the schema
            if isinstance(schema, dict):
                # Simple field definitions: {"field_name": "type"}
                field_definitions = {}
                type_mapping = {
                    "str": str,
                    "string": str,
                    "int": int,
                    "integer": int,
                    "float": float,
                    "number": float,
                    "bool": bool,
                    "boolean": bool,
                    "list": list,
                    "array": list,
                    "dict": dict,
                    "object": dict,
                }

                for field_name, field_type in schema.items():
                    if isinstance(field_type, str):
                        py_type = type_mapping.get(field_type.lower(), Any)
                        field_definitions[field_name] = (py_type, ...)
                    else:
                        field_definitions[field_name] = (Any, ...)

                if field_definitions:
                    return create_model("TaskOutputModel", **field_definitions)
            else:
                self.log(
                    f"Invalid JSON schema format in task '{self.task_name}': "
                    f"expected dictionary, got {type(schema).__name__}"
                )

        except Exception as e:  # noqa: BLE001
            self.log(f"Error parsing JSON schema in task '{self.task_name}': {e}")

        return None

    def build_task(self) -> Any:
        """Build and return the PraisonAI Task instance."""
        from praisonai.flow.helpers import convert_tools

        task_class = self._import_task()

        # Convert tools if provided
        tools = convert_tools(self.tools) if self.tools else None

        # Build context from connected tasks
        context = None
        if self.context:
            context = [t for t in self.context if t is not None]
            if not context:
                context = None

        # Parse output JSON schema
        output_json = self._parse_output_json()

        # Parse next_tasks
        next_tasks = None
        if self.next_tasks:
            next_tasks = [t.strip() for t in self.next_tasks.split(",") if t.strip()]

        # Parse images
        images = None
        if self.images:
            images = [img for img in self.images if img is not None]
            if not images:
                images = None

        # Build task kwargs
        kwargs = {
            "name": self.task_name,
            "description": self.description,
            "expected_output": self.expected_output,
            "agent": self.agent,
            "tools": tools if tools else [],
            "context": context if context else [],
            "async_execution": self.async_execution,
            "quality_check": self.quality_check,
            "max_retries": self.max_retries,
            "rerun": self.rerun,
            "retain_full_context": self.retain_full_context,
        }

        # Add optional parameters
        if output_json:
            kwargs["output_json"] = output_json

        if self.output_file:
            kwargs["output_file"] = self.output_file
            kwargs["create_directory"] = self.create_directory

        if self.input_file:
            kwargs["input_file"] = self.input_file

        if images:
            kwargs["images"] = images

        if self.guardrail:
            kwargs["guardrail"] = True

        if self.task_type and self.task_type != "task":
            kwargs["task_type"] = self.task_type

        if self.condition:
            kwargs["condition"] = self.condition

        if next_tasks:
            kwargs["next_tasks"] = next_tasks

        if self.is_start:
            kwargs["is_start"] = True

        if self.variables:
            kwargs["variables"] = self.variables

        # Build task
        task = task_class(**kwargs)

        self.status = f"Task '{self.task_name}' created"
        return task
