"""Agent runtime for Gremlin Autopilot.

Implements the LLM agent that:
1. Accepts chat messages with sheet context
2. Calls LLM with tool definitions
3. Parses tool calls and generates preview
4. Returns plans for user approval
5. Executes approved plans
"""

from __future__ import annotations

import hashlib
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from .models import (
    ActionResult,
    ChatRequest,
    ChatResponse,
    ExecuteResponse,
    JobStatus,
    PlanAction,
    SheetContext,
    UserSettings,
)
from .security import build_llm_context, validate_tool_args
from .tools import get_registry_for_user

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)

# System prompt for the agent
SYSTEM_PROMPT = """You are Gremlin, a helpful data operations assistant for RevOps teams.

Your job is to help users clean, transform, enrich, and manage their data in Google Sheets.

You have access to a set of tools that can operate on spreadsheet data. When the user asks
for help, analyze their request and call the appropriate tools.

IMPORTANT RULES:
1. Only use the tools provided - never suggest manual steps or code.
2. For each tool call, specify all required parameters.
3. Prefer simple, reversible operations when possible.
4. If you're unsure which column to use, ask the user to clarify.
5. Never make up column names - use only the columns from the provided schema.

When you receive sheet context, you'll see:
- headers: List of column names in the sheet
- row_count: Number of data rows
- column_types: Inferred types for each column (email, phone, company, etc.)

Use this information to select appropriate tools and parameters.
"""


@dataclass
class Plan:
    """Represents a plan of actions to execute."""

    plan_id: str
    context_hash: str
    actions: list[PlanAction]
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    status: str = "pending"
    message: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Convert plan to dictionary."""
        return {
            "plan_id": self.plan_id,
            "context_hash": self.context_hash,
            "actions": [a.model_dump() for a in self.actions],
            "created_at": self.created_at.isoformat(),
            "status": self.status,
            "message": self.message,
        }


class StaleContextError(Exception):
    """Raised when sheet context has changed since plan creation."""

    pass


class AgentRuntime:
    """Core agent runtime for processing chat and executing tools."""

    def __init__(
        self,
        api_key: str | None = None,
        provider: str = "anthropic",
        model: str = "claude-sonnet-4-20250514",
        user_credentials: set[str] | None = None,
        user_settings: UserSettings | None = None,
    ):
        """Initialize the agent runtime.

        Args:
            api_key: LLM API key. If None, agent runs in "template-only" mode.
            provider: LLM provider (anthropic, openai).
            model: Model identifier to use.
            user_credentials: Set of credential types the user has configured.
            user_settings: User settings for consent and preferences.
        """
        self.api_key = api_key
        self.provider = provider
        self.model = model
        self.user_credentials = user_credentials or set()
        self.user_settings = user_settings or UserSettings()
        self.tool_registry = get_registry_for_user(self.user_credentials)

        # In-memory plan store (would be SQLite in production)
        self._plans: dict[str, Plan] = {}

        # Initialize LLM client if API key provided
        self._client = None
        if api_key:
            self._init_llm_client()

    def _init_llm_client(self) -> None:
        """Initialize the LLM client based on provider."""
        if self.provider == "anthropic":
            try:
                from anthropic import Anthropic

                self._client = Anthropic(api_key=self.api_key)
            except ImportError:
                logger.warning("anthropic package not installed; agent mode disabled")
        elif self.provider == "openai":
            try:
                from openai import OpenAI

                self._client = OpenAI(api_key=self.api_key)
            except ImportError:
                logger.warning("openai package not installed; agent mode disabled")

    def _compute_context_hash(self, context: SheetContext) -> str:
        """Compute a hash of the sheet context for staleness detection."""
        payload = f"{context.sheet_id}:{context.tab_name}:{context.row_count}"
        if context.headers:
            payload += ":" + ",".join(context.headers)
        return hashlib.sha256(payload.encode()).hexdigest()[:16]

    def _get_tools_for_llm(self) -> list[dict[str, Any]]:
        """Get tool definitions formatted for LLM."""
        return self.tool_registry.get_tools_for_llm()

    async def process_chat(
        self,
        request: ChatRequest,
        sheet_context: SheetContext | None = None,
    ) -> ChatResponse:
        """Process a chat message and generate an action plan.

        Args:
            request: The chat request with user message.
            sheet_context: Optional context about the active sheet.

        Returns:
            ChatResponse with plan or error message.
        """
        if not self._client:
            return ChatResponse(
                status=JobStatus.FAILED,
                message="LLM not configured. Please add an API key in settings.",
                job_id=None,
            )

        try:
            # Build LLM context (respecting consent settings)
            llm_context = None
            if sheet_context:
                llm_context = build_llm_context(sheet_context, self.user_settings)

            # Build messages
            messages = self._build_messages(request.message, llm_context)

            # Get tool definitions
            tools = self._get_tools_for_llm()

            # Call LLM
            if self.provider == "anthropic":
                response = await self._call_anthropic(messages, tools)
            else:
                response = await self._call_openai(messages, tools)

            # Parse tool calls from response
            actions = self._parse_tool_calls(response)

            if not actions:
                # No tool calls - just a text response
                return ChatResponse(
                    status=JobStatus.COMPLETED,
                    message=response.get("content", "I'm not sure how to help with that."),
                    job_id=None,
                )

            # Create plan
            context_hash = ""
            if sheet_context:
                context_hash = self._compute_context_hash(sheet_context)

            plan = Plan(
                plan_id=f"plan_{uuid4().hex[:8]}",
                context_hash=context_hash,
                actions=actions,
                message=f"Generated {len(actions)} action(s) for approval.",
            )

            # Store plan
            self._plans[plan.plan_id] = plan

            return ChatResponse(
                status=JobStatus.PLAN_READY,
                message=plan.message,
                job_id=plan.plan_id,
                actions=actions,
            )

        except Exception as e:
            logger.exception("Error processing chat")
            return ChatResponse(
                status=JobStatus.FAILED,
                message=f"Error: {e!s}",
                job_id=None,
            )

    def _build_messages(
        self,
        user_message: str,
        llm_context: Any | None = None,
    ) -> list[dict[str, str]]:
        """Build message list for LLM."""
        messages = []

        # Add system prompt
        system_content = SYSTEM_PROMPT

        # Add sheet context if available
        if llm_context:
            context_str = self._format_context_for_llm(llm_context)
            system_content += f"\n\n## Current Sheet Context\n{context_str}"

        messages.append({"role": "system", "content": system_content})
        messages.append({"role": "user", "content": user_message})

        return messages

    def _format_context_for_llm(self, context: Any) -> str:
        """Format sheet context for LLM prompt."""
        lines = [
            f"Sheet ID: {context.sheet_id}",
            f"Tab: {context.tab_name}",
            f"Row count: {context.row_count}",
            f"Headers: {', '.join(context.headers)}",
        ]

        if hasattr(context, "column_types") and context.column_types:
            lines.append("Column types:")
            for col, col_type in context.column_types.items():
                lines.append(f"  - {col}: {col_type}")

        if hasattr(context, "sample_rows") and context.sample_rows:
            lines.append("\nSample rows:")
            for i, row in enumerate(context.sample_rows[:3]):
                lines.append(f"  Row {i + 1}: {json.dumps(row)}")

        return "\n".join(lines)

    async def _call_anthropic(
        self,
        messages: list[dict[str, str]],
        tools: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """Call Anthropic API with tool use."""
        from anthropic import AsyncAnthropic

        # Use async client
        client = AsyncAnthropic(api_key=self.api_key)

        # Convert tools to Anthropic format
        anthropic_tools = []
        for tool in tools:
            anthropic_tools.append({
                "name": tool["name"],
                "description": tool["description"],
                "input_schema": tool["parameters"],
            })

        # Extract system message
        system_msg = next(
            (m["content"] for m in messages if m["role"] == "system"),
            "",
        )
        user_messages = [m for m in messages if m["role"] != "system"]

        response = await client.messages.create(
            model=self.model,
            max_tokens=4096,
            system=system_msg,
            messages=user_messages,
            tools=anthropic_tools,
        )

        # Parse response
        result: dict[str, Any] = {"tool_calls": [], "content": ""}

        for block in response.content:
            if block.type == "text":
                result["content"] = block.text
            elif block.type == "tool_use":
                result["tool_calls"].append({
                    "name": block.name,
                    "args": block.input,
                })

        return result

    async def _call_openai(
        self,
        messages: list[dict[str, str]],
        tools: list[dict[str, Any]],
    ) -> dict[str, Any]:
        """Call OpenAI API with tool use."""
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=self.api_key)

        # Convert tools to OpenAI format
        openai_tools = []
        for tool in tools:
            openai_tools.append({
                "type": "function",
                "function": {
                    "name": tool["name"],
                    "description": tool["description"],
                    "parameters": tool["parameters"],
                },
            })

        response = await client.chat.completions.create(
            model=self.model if "gpt" in self.model else "gpt-4o",
            messages=messages,
            tools=openai_tools if openai_tools else None,
        )

        # Parse response
        result: dict[str, Any] = {"tool_calls": [], "content": ""}
        choice = response.choices[0]

        if choice.message.content:
            result["content"] = choice.message.content

        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                result["tool_calls"].append({
                    "name": tc.function.name,
                    "args": json.loads(tc.function.arguments),
                })

        return result

    def _parse_tool_calls(self, response: dict[str, Any]) -> list[PlanAction]:
        """Parse tool calls from LLM response into PlanActions."""
        actions = []

        for tool_call in response.get("tool_calls", []):
            tool_name = tool_call["name"]
            tool_args = tool_call["args"]

            # Validate tool exists
            tool_def = self.tool_registry.get_tool(tool_name)
            if not tool_def:
                logger.warning(f"Unknown tool: {tool_name}")
                continue

            # Validate and sanitize arguments
            try:
                validate_tool_args(tool_name, tool_args)
            except Exception as e:
                logger.warning(f"Invalid args for {tool_name}: {e}")
                continue

            # Create action
            action = PlanAction(
                tool=tool_name,
                tool_ref=tool_def.name,
                args=tool_args,
                label=f"{tool_def.description}",
                risk_level=int(tool_def.risk_level) if tool_def.risk_level is not None else 3,
                destructive=tool_def.destructive,
                writes_external=tool_def.writes_external,
                rollback_supported=tool_def.rollback_supported,
                preview_trustworthy=tool_def.preview_trustworthy,
                cli_equivalent=self._format_cli_command(tool_def, tool_args),
            )
            actions.append(action)

        return actions

    def _format_cli_command(self, tool_def: Any, args: dict[str, Any]) -> str:
        """Format CLI equivalent command for display."""
        if not tool_def.cli_equivalent:
            return ""

        cmd = tool_def.cli_equivalent
        for key, value in args.items():
            placeholder = "{" + key + "}"
            if placeholder in cmd:
                cmd = cmd.replace(placeholder, str(value))
        return cmd

    async def execute_plan(
        self,
        plan_id: str,
        current_context: SheetContext | None = None,
    ) -> ExecuteResponse:
        """Execute an approved plan.

        Args:
            plan_id: ID of the plan to execute.
            current_context: Current sheet context for staleness check.

        Returns:
            ExecuteResponse with results.
        """
        plan = self._plans.get(plan_id)
        if not plan:
            return ExecuteResponse(
                status=JobStatus.FAILED,
                job_id=plan_id,
                results=[],
                message="Plan not found",
            )

        # Check for stale context
        if current_context and plan.context_hash:
            current_hash = self._compute_context_hash(current_context)
            if current_hash != plan.context_hash:
                return ExecuteResponse(
                    status=JobStatus.FAILED,
                    job_id=plan_id,
                    results=[],
                    message="Sheet has changed since plan was created. Please regenerate.",
                )

        # Execute each action
        results: list[ActionResult] = []
        for action in plan.actions:
            try:
                result = await self._execute_action(action)
                results.append(result)

                if result.status == "error":
                    # Stop on first error
                    plan.status = "failed"
                    return ExecuteResponse(
                        status=JobStatus.FAILED,
                        job_id=plan_id,
                        results=results,
                        message=f"Failed at action: {action.tool}",
                    )

            except Exception as e:
                logger.exception(f"Error executing {action.tool}")
                results.append(ActionResult(
                    tool=action.tool,
                    status="error",
                    rows_affected=0,
                    message=str(e),
                ))
                plan.status = "failed"
                return ExecuteResponse(
                    status=JobStatus.FAILED,
                    job_id=plan_id,
                    results=results,
                    message=str(e),
                )

        plan.status = "completed"
        return ExecuteResponse(
            status=JobStatus.COMPLETED,
            job_id=plan_id,
            results=results,
            message=f"Successfully executed {len(results)} action(s)",
        )

    async def _execute_action(self, action: PlanAction) -> ActionResult:
        """Execute a single action.

        Agent-based tool execution is not yet implemented. Use the CLI
        play runner instead:  g-gremlin autopilot play run <name>
        """
        logger.warning(
            "Agent tool execution is not yet available. "
            "Use 'g-gremlin autopilot play run' for workflow execution."
        )
        return ActionResult(
            tool=action.tool,
            status="error",
            rows_affected=0,
            message=(
                "Agent execution is not yet available. "
                "Use the CLI instead: g-gremlin autopilot play run <play_name>"
            ),
        )

    def get_plan(self, plan_id: str) -> Plan | None:
        """Get a plan by ID."""
        return self._plans.get(plan_id)


def create_agent(
    api_key: str | None = None,
    provider: str = "anthropic",
    user_credentials: set[str] | None = None,
) -> AgentRuntime:
    """Factory function to create an agent runtime.

    Args:
        api_key: LLM API key.
        provider: LLM provider name.
        user_credentials: Set of configured credential types.

    Returns:
        Configured AgentRuntime instance.
    """
    return AgentRuntime(
        api_key=api_key,
        provider=provider,
        user_credentials=user_credentials,
    )
