"""Data processing tool for generating SQL feature engineering YAML configurations."""

from __future__ import annotations

import json
import uuid
from typing import TYPE_CHECKING

import yaml

from arc.tools.base import BaseTool, ToolResult
from arc.utils import ConfirmationService
from arc.utils.yaml_workflow import YamlConfirmationWorkflow

if TYPE_CHECKING:
    from arc.core.agents.ml_data.ml_data import (
        MLDataAgent,
    )
    from arc.database.services.container import ServiceContainer
else:
    # Avoid circular imports at runtime
    try:
        from arc.core.agents.ml_data.ml_data import (
            MLDataAgent,
            MLDataError,
        )
    except ImportError:
        MLDataAgent = None
        MLDataError = Exception


class MLDataTool(BaseTool):
    """Tool for generating data processing YAML configurations from natural language.

    This tool has a single purpose: convert natural language descriptions into
    structured SQL data processing pipelines using LLM generation.
    It requires an API key and uses the MLDataAgent.
    """

    def __init__(
        self,
        services: ServiceContainer,
        api_key: str,
        base_url: str | None = None,
        model: str | None = None,
        ui_interface=None,
    ):
        """Initialize MLDataTool.

        Args:
            services: ServiceContainer instance providing database access
            api_key: API key for LLM calls (can be empty for enterprise gateway)
            base_url: Base URL for LLM calls
            model: Model name for LLM calls
            ui_interface: UI interface for interactive confirmation workflow

        Raises:
            RuntimeError: If MLDataAgent cannot be initialized

        Note:
            Empty API keys are allowed to support enterprise gateway environments
            where authentication is handled by the gateway (only base_url needed).
            The validation is handled at the CLI level in _get_ml_tool_config().
        """
        # Allow empty API key for enterprise gateway mode
        # The OpenAI SDK will use the provided base_url for authentication

        if MLDataAgent is None:
            raise RuntimeError(
                "MLDataAgent is not available. "
                "This is likely due to missing dependencies or circular imports."
            )

        self.services = services
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.ui = ui_interface

        # Initialize the generator agent (required)
        try:
            self.generator_agent = MLDataAgent(
                services=services, api_key=api_key, base_url=base_url, model=model
            )
        except Exception as e:
            raise RuntimeError(
                f"Failed to initialize MLDataAgent: {str(e)}. "
                "This tool requires a working LLM connection for data processing "
                "generation."
            ) from e

    def _parse_sql_to_array(self, sql_string: str) -> list[str]:
        """Parse SQL string into array of clean SQL statements.

        Splits the SQL by step markers and removes transaction boilerplate.

        Args:
            sql_string: Full SQL string with step markers and transaction

        Returns:
            List of clean SQL statements without step markers
        """
        if not sql_string:
            return []

        import re

        # Remove transaction markers from the entire SQL first
        # This handles cases where COMMIT appears at the end of SQL blocks
        sql_string = sql_string.replace("BEGIN TRANSACTION;", "")
        sql_string = sql_string.replace("COMMIT;", "")

        # Split by step markers (e.g., "-- Step 1/6: ...")
        # Pattern: -- Step X/Y: step_name (type)
        step_pattern = r"-- Step \d+/\d+: [^\n]+\n"

        # Split on step markers but keep the SQL parts
        parts = re.split(step_pattern, sql_string)

        # Clean and collect non-empty SQL statements
        sql_statements = []
        for part in parts:
            # Clean up whitespace
            cleaned = part.strip()

            # Skip empty parts
            if not cleaned:
                continue

            sql_statements.append(cleaned)

        return sql_statements

    def _build_data_result(
        self,
        status: str,
        data_processing_id: str,
        output_tables: list[str] | None = None,
        sql_operations: list[str] | None = None,
    ) -> str:
        """Build structured JSON result for ML data tool.

        Args:
            status: "accepted" or "cancelled"
            data_processing_id: Data processing execution ID
            output_tables: List of created table names
            sql_operations: List of SQL statements (or single string to parse)

        Returns:
            JSON string with structured data processing result
        """
        # Parse SQL if it's a single string
        if sql_operations and isinstance(sql_operations, str):
            sql_operations = self._parse_sql_to_array(sql_operations)

        result = {
            "status": status,
            "data_processing_id": data_processing_id,
            "execution": {
                "output_tables": output_tables or [],
                "sql_operations": sql_operations or [],
            },
        }
        return json.dumps(result)

    async def generate(
        self,
        name: str,
        data_source_type: str,
        data_sources: list[str],
        instruction: str | None = None,
        database: str = "user",
        auto_confirm: bool = False,
        knowledge_references: list[str] | None = None,
    ) -> ToolResult:
        """Generate YAML configuration from instruction using LLM.

        Args:
            name: Name for the data processor (will be registered in database)
            data_source_type: Type of data sources - "csv", "parquet", "json", or "table"
            data_sources: List of data sources (file paths/URLs or table names)
            instruction: Detailed instruction for data processing
            database: Database to use - "system" or "user"
            auto_confirm: Skip interactive confirmation workflow
            knowledge_references: Optional list of knowledge IDs referenced by
                this request (e.g., ['ml_data_preparation'])

        Returns:
            ToolResult with operation result
        """
        # Build metadata for section title (do this before validation for clean code)
        metadata_parts = []
        if knowledge_references:
            metadata_parts.extend(knowledge_references)

        # Use context manager for section printing (automatic cleanup and spacing)
        with self._section_printer(
            self.ui, "ML Data", metadata=metadata_parts
        ) as printer:
            # Helper to show error and return (define early for use in validation)
            def _error_in_section(message: str) -> ToolResult:
                if printer:
                    printer.print("")
                    printer.print(f"✗ {message}")
                return ToolResult(
                    success=False,
                    output=message,
                    metadata={"error_shown": True, "error_message": message},
                )

            # Validate required parameters
            if not name:
                return _error_in_section(
                    "Name is required for data processor registration. "
                    "Provide a name for the data processor."
                )

            if not instruction:
                return _error_in_section(
                    "Instruction is required for YAML generation. "
                    "Provide a detailed instruction for your data processing needs."
                )

            # Validate data_source_type
            if not data_source_type:
                return _error_in_section(
                    "data_source_type is required. "
                    "Specify 'csv', 'parquet', 'json', or 'table'."
                )

            valid_types = ["csv", "parquet", "json", "table"]
            if data_source_type not in valid_types:
                return _error_in_section(
                    f"Invalid data_source_type: {data_source_type}. "
                    f"Must be one of: {', '.join(valid_types)}."
                )

            # Validate data_sources with type-specific hints
            if not data_sources or len(data_sources) == 0:
                hints = {
                    "csv": "CSV file paths or URLs (e.g., ['data.csv', 'https://example.com/data.csv'])",
                    "parquet": "Parquet file paths or URLs (e.g., ['data.parquet', 'https://...'])",
                    "json": "JSON file paths or URLs (e.g., ['data.json', 'https://...'])",
                    "table": "table names (e.g., ['users', 'transactions'])",
                }
                hint = hints.get(data_source_type, "data sources")

                return _error_in_section(
                    "data_sources is required to narrow the scope of data exploration. "
                    f"For data_source_type='{data_source_type}', specify {hint}."
                )

            # Validate database
            if database not in ["system", "user"]:
                return _error_in_section(
                    f"Invalid database: {database}. Must be 'system' or 'user'."
                )

            # Use instruction directly (no plan loading needed)
            enhanced_instruction = instruction

            # Show task description only in verbose mode
            if printer:
                # Check verbose mode from settings
                from arc.core.config import SettingsManager

                settings = SettingsManager()
                if settings.get_verbose_mode():
                    printer.print(f"[dim]Task: {instruction}[/dim]")
                    printer.print("")  # Empty line after task

            try:
                # Set progress callback for this invocation
                # (agent can be reused with different printers)
                if printer:
                    self.generator_agent.progress_callback = printer.print
                else:
                    self.generator_agent.progress_callback = None

                # Generate using LLM (generator_agent is guaranteed to exist)
                (
                    spec,
                    yaml_content,
                    conversation_history,  # Store for interactive editing workflow
                ) = await self.generator_agent.generate_data_processing_yaml(
                    instruction=enhanced_instruction,
                    name=name,
                    data_source_type=data_source_type,
                    data_sources=data_sources,
                    database=database,
                    knowledge_references=knowledge_references,
                )

                # Show completion message
                if printer:
                    printer.print("[dim]✓ Data processor generated[/dim]")

                # Spec is already validated by agent (with retries)
                # No need for redundant validation here

                # Interactive confirmation workflow
                # (unless auto_confirm is True or no UI available)
                # Get UI from singleton if not provided directly
                ui = self.ui
                if ui is None:
                    ui = ConfirmationService.get_instance()._ui

                if not auto_confirm and ui:
                    workflow = YamlConfirmationWorkflow(
                        validator_func=self._create_validator(),
                        editor_func=self._create_editor(),
                        ui_interface=ui,
                        yaml_type_name="data processor",
                        yaml_suffix=".arc-data.yaml",
                    )

                    context_dict = {
                        "name": name,
                        "instruction": str(enhanced_instruction),
                        "data_source_type": data_source_type,
                        "data_sources": data_sources,
                        "database": database,
                    }

                    try:
                        proceed, final_yaml = await workflow.run_workflow(
                            yaml_content,
                            context_dict,
                            None,  # No output path - saved to DB only
                            conversation_history,  # Pass history for editing
                        )
                        if not proceed:
                            # Generate placeholder data_processing_id for cancelled state
                            cancelled_id = f"cancelled_{uuid.uuid4().hex[:8]}"
                            output_json = self._build_data_result(
                                status="cancelled",
                                data_processing_id=cancelled_id,
                                output_tables=[],
                                sql_operations=[],
                            )
                            return ToolResult(
                                success=True,
                                output=output_json,
                                metadata={"cancelled": True},
                            )
                        yaml_content = final_yaml
                        # Re-parse spec from edited YAML
                        from arc.graph.features.data_source import DataSourceSpec

                        spec = DataSourceSpec.from_yaml(yaml_content)
                    finally:
                        workflow.cleanup()

                # Register data processor in database
                from arc.ml.runtime import MLRuntime, MLRuntimeError

                runtime = MLRuntime(self.services, artifacts_dir="artifacts")
                try:
                    # Use spec.description if available, otherwise use name
                    description = spec.description if spec.description else name
                    processor = runtime.register_data_processor(
                        name=name, spec=spec, description=description
                    )
                    # Display registration confirmation in the Data Processor section
                    if printer:
                        printer.print("")
                        printer.print(
                            f"[dim]✓ Data processor registered: {processor.id}[/dim]"
                        )
                        printer.print(f"[dim]  {len(spec.steps)} steps[/dim]")
                except MLRuntimeError as e:
                    return _error_in_section(
                        f"Failed to register data processor: {str(e)}"
                    )

                # Execute the pipeline automatically after confirmation
                from arc.ml.data_source_executor import (
                    DataSourceExecutionError,
                    execute_data_source_pipeline,
                )

                # Show execution message
                if printer:
                    printer.print("")
                    printer.print("→ Executing data processing pipeline")

                # Define progress callback for real-time updates
                def progress_callback(message: str, level: str):
                    """Handle progress updates during execution."""
                    if printer:
                        if level == "step":
                            printer.print(f"[dim]  {message}[/dim]")
                        elif level == "warning":
                            printer.print(f"[yellow]  ⚠️ {message}[/yellow]")
                        elif level == "error":
                            printer.print(f"[red]  ❌ {message}[/red]")

                try:
                    execution_result = await execute_data_source_pipeline(
                        spec, database, self.services.db_manager, progress_callback
                    )

                    # Invalidate schema cache since new tables were created
                    self.services.schema.invalidate_cache(database)

                    # Use processor.id as data_processing_id for consistency with other ML tools
                    # (e.g., "diabetes-feature-processing-v3" instead of random "data_98f4ac51")
                    data_processing_id = processor.id

                    # Note: Execution records are no longer stored since we removed ml_plan
                    # The data processor registration already tracks all necessary metadata

                except DataSourceExecutionError as e:
                    # Generation and registration succeeded, but execution failed
                    if printer:
                        printer.print("")
                        printer.print(f"⚠ Pipeline execution failed: {str(e)}")
                        printer.print("")
                        printer.print(
                            "[dim]The data processor was successfully generated and "
                            "registered to the database.[/dim]"
                        )
                        printer.print(
                            "[dim]You can review and fix the SQL queries, then re-run "
                            "the processor.[/dim]"
                        )

                    # Return FAILURE with detailed metadata for agent to understand
                    # Generation/registration succeeded, but execution failed
                    return ToolResult(
                        success=False,  # Execution failed - agent needs to know
                        output=f"Data processor '{name}' registered as {processor.id}, "
                        f"but execution failed: {str(e)}",
                        metadata={
                            "processor_id": processor.id,
                            "processor_name": name,
                            "generation_succeeded": True,
                            "registration_succeeded": True,
                            "execution_succeeded": False,
                            "execution_error": str(e),
                            "error_type": "execution_failure",
                        },
                    )

                # Show success summary
                if printer:
                    printer.print("")
                    printer.print("[dim]✓ Pipeline executed successfully[/dim]")
                    printer.print(
                        f"[dim]  Table: {', '.join(execution_result.created_tables)}[/dim]"
                    )
                    printer.print(
                        f"[dim]  Time: {execution_result.execution_time:.2f}s[/dim]"
                    )

                # Build structured JSON output
                output_json = self._build_data_result(
                    status="accepted",
                    data_processing_id=data_processing_id,
                    output_tables=execution_result.created_tables,
                    sql_operations=execution_result.sql or [],
                )

                # Build metadata
                metadata = {
                    "processor_id": processor.id,
                    "processor_name": name,
                    "generation_succeeded": True,
                    "registration_succeeded": True,
                    "execution_succeeded": True,
                    "created_tables": execution_result.created_tables,
                    "execution_time": execution_result.execution_time,
                    "data_processing_id": data_processing_id,
                }

                # Return success with structured JSON output
                return ToolResult(
                    success=True,
                    output=output_json,
                    metadata=metadata,
                )

            except Exception as e:
                # Preserve validation error details for main LLM
                error_msg = str(e)
                return _error_in_section(
                    f"Failed to generate data processor after 3 attempts: {error_msg}"
                )

    async def execute(self, **kwargs) -> ToolResult:
        """Execute method for BaseTool compatibility."""
        return await self.generate(**kwargs)

    def _create_validator(self):
        """Create validator function for the workflow.

        Returns:
            Validator function that takes YAML string and returns list of errors
        """
        from arc.graph.features.data_source import DataSourceSpec

        def validate(yaml_str: str) -> list[str]:
            """Validate YAML string as DataSourceSpec.

            Args:
                yaml_str: YAML string to validate

            Returns:
                List of validation error messages (empty if valid)
            """
            try:
                # Try to parse and validate as DataSourceSpec
                DataSourceSpec.from_yaml(yaml_str)
                return []  # No errors
            except yaml.YAMLError as e:
                return [f"Invalid YAML: {e}"]
            except ValueError as e:
                return [f"Validation error: {e}"]
            except Exception as e:
                return [f"Unexpected error: {e}"]

        return validate

    def _create_editor(self):
        """Create editor function for AI-assisted editing in the workflow.

        Returns:
            Editor function that takes (yaml, feedback, context, conversation_history)
            and returns tuple of (edited_yaml, updated_history)
        """

        async def edit(
            yaml_content: str,
            feedback: str,
            context: dict[str, str],
            conversation_history: list[dict[str, str]] | None = None,
        ) -> tuple[str | None, list[dict[str, str]] | None]:
            """Edit YAML with AI assistance using conversation history.

            Args:
                yaml_content: Current YAML content
                feedback: User feedback describing desired changes
                context: Context dictionary with generation parameters
                conversation_history: Optional conversation history for
                    continuing conversation

            Returns:
                Tuple of (edited_yaml, updated_conversation_history) or
                (None, None) if failed
            """
            try:
                # Set progress callback for editing invocation
                # Note: progress_callback should already be set from generate(),
                # but we keep it here for robustness
                # (it's the same agent instance, so callback is preserved)

                # Call the generator agent with conversation history
                # feedback becomes the instruction in edit mode
                (
                    _spec,
                    edited_yaml,
                    updated_history,
                ) = await self.generator_agent.generate_data_processing_yaml(
                    instruction=feedback,  # User's change request
                    name=context.get("name", ""),
                    data_source_type=context.get("data_source_type", "table"),
                    data_sources=context.get("data_sources", []),
                    database=context.get("database", "user"),
                    existing_yaml=yaml_content,
                    knowledge_references=None,  # Editing uses conversation_history
                    conversation_history=conversation_history,  # Continue conversation
                )
                return edited_yaml, updated_history

            except Exception as e:
                if self.ui:
                    self.ui.show_system_error(f"❌ AI editing failed: {str(e)}")
                return None, None

        return edit
