"""Cortex COMPLETE/SUMMARIZE wrapper agent.

Wraps the CortexClient (src/cortex_agent/cortex_client.py) for
LLM-powered text generation, summarization, and sentiment analysis
via Snowflake Cortex.
"""

from __future__ import annotations

import logging
from typing import Any, Dict, Optional

from .base_agent import BaseDataBridgeAgent
from .state_schema import WorkflowState
from .databridge_tools import tool_cortex_complete

logger = logging.getLogger(__name__)


class CortexLLMAgent(BaseDataBridgeAgent):
    """Specialist agent for Snowflake Cortex LLM operations."""

    name = "cortex_llm_agent"
    description = "Runs Snowflake Cortex COMPLETE, SUMMARIZE, and SENTIMENT functions"
    phase_name = "cortex_llm"

    def __init__(self, **kwargs):
        super().__init__(
            tools=[tool_cortex_complete],
            system_prompt=(
                "You are the Cortex LLM agent for DataBridge AI. "
                "You use Snowflake Cortex functions (COMPLETE, SUMMARIZE, SENTIMENT) "
                "to perform AI operations within the Snowflake environment."
            ),
            **kwargs,
        )
        self._cortex_client = None

    def set_cortex_client(self, client) -> None:
        """Inject a configured CortexClient instance."""
        self._cortex_client = client

    async def run(self, state: WorkflowState) -> WorkflowState:
        """Execute Cortex LLM operation."""
        config = state.get("config", {})
        action = config.get("cortex_action", "complete")

        if not self._cortex_client:
            return self._update_context(state, self.phase_name, {
                "status": "not_configured",
                "note": "CortexClient not injected — requires Snowflake connection",
            })

        if action == "complete":
            return await self._run_complete(state, config)
        elif action == "summarize":
            return await self._run_summarize(state, config)
        elif action == "sentiment":
            return await self._run_sentiment(state, config)

        return self._update_context(state, self.phase_name, {
            "error": f"Unknown cortex action: {action}"
        })

    async def _run_complete(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Run Cortex COMPLETE."""
        prompt = config.get("cortex_prompt", "")
        if not prompt:
            return self._update_context(state, self.phase_name, {
                "error": "No prompt provided for COMPLETE"
            })

        try:
            result = self._cortex_client.complete(
                prompt=prompt,
                model=config.get("cortex_model", "mistral-large"),
                temperature=config.get("cortex_temperature", 0.3),
            )
            return self._update_context(state, self.phase_name, {
                "action": "complete",
                "result": result.text if hasattr(result, "text") else str(result),
            })
        except Exception as e:
            return self._update_context(state, self.phase_name, {"error": str(e)})

    async def _run_summarize(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Run Cortex SUMMARIZE."""
        text = config.get("cortex_text", "")
        if not text:
            return self._update_context(state, self.phase_name, {
                "error": "No text provided for SUMMARIZE"
            })

        try:
            result = self._cortex_client.summarize(text=text)
            return self._update_context(state, self.phase_name, {
                "action": "summarize",
                "result": result.text if hasattr(result, "text") else str(result),
            })
        except Exception as e:
            return self._update_context(state, self.phase_name, {"error": str(e)})

    async def _run_sentiment(self, state: WorkflowState, config: Dict) -> WorkflowState:
        """Run Cortex SENTIMENT."""
        text = config.get("cortex_text", "")
        if not text:
            return self._update_context(state, self.phase_name, {
                "error": "No text provided for SENTIMENT"
            })

        try:
            result = self._cortex_client.sentiment(text=text)
            return self._update_context(state, self.phase_name, {
                "action": "sentiment",
                "result": result.text if hasattr(result, "text") else str(result),
            })
        except Exception as e:
            return self._update_context(state, self.phase_name, {"error": str(e)})
