from __future__ import annotations

import asyncio
import datetime
import logging
import uuid
from abc import ABC, abstractmethod
from collections.abc import Callable
from typing import Any

import langsmith as ls
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel

from broadie.config import settings

from .channels import ChannelsAgent
from .mixins import PersistenceMixin
from .prompts import BASE_INSTRUCTIONS
from .schemas import AgentSchema, BaseDefaultOutput, ChannelSchema, ModelSchema, SubAgentSchema
from .tools.tasks import create_tasks, update_task
from .utils import slugify

logger = logging.getLogger(__name__)


class BaseAgent(PersistenceMixin, ABC):
    def __init__(
        self,
        name: str,
        model: ModelSchema,
        tools: list[str | Callable],
        strict_mode: bool = True,
        label: str | None = None,
        instruction: str = "You are a helpful assistant.",
        output_schema: dict[str, Any] | None = None,
        deliver_to_channels: bool = True,
        history_turns: int = 10,  # Number of previous conversation turns to include in context
        conversational: bool = True,  # Whether this agent uses conversation history
    ):
        PersistenceMixin.__init__(self)
        super().__init__()
        self.label = label or name
        self.id = slugify(name)
        self.name = name
        self.model = model
        self.tools = tools
        self.instruction = instruction
        self.strict_mode = strict_mode
        self.output_schema = output_schema or BaseDefaultOutput
        self.runtime = None
        self.deliver_to_channels = deliver_to_channels
        self.channels = []
        self.channels_agent = None
        self.history_turns = history_turns  # Store history configuration
        self.conversational = conversational  # Control conversation history usage

        # Create model instance and structured LLM
        self.model_instance = None
        self.structured_llm = None

    @staticmethod
    def _inbuilt_tools() -> list[Callable]:
        """Always available tools for all agents."""
        return [create_tasks, update_task]

    @property
    def recursion_limit(self) -> int:
        return settings.RECURSION_LIMIT or max(
            15,
            len(self._inbuilt_tools() + self.agent_tools()) + 10,
        )

    def initialize_runtime(self):
        """Initialize direct LLM approach with proper tool binding and structured output."""
        self._initialize_direct_llm()

    def _initialize_direct_llm(self):
        """Initialize direct LLM approach with proper tool binding and structured output."""
        # Initialize the appropriate model instance
        model_config = {
            "temperature": 0.2,
            "max_tokens": 8100,
        }

        if self.model.provider.value == "google_vertexai":
            try:
                from langchain_google_vertexai import ChatVertexAI

                self.model_instance = ChatVertexAI(model_name=self.model.name, **model_config)
            except ImportError as e:
                raise ImportError(
                    "Google Vertex AI provider requires 'langchain-google-vertexai'. "
                    "Install it with: pip install langchain-google-vertexai"
                ) from e

        elif self.model.provider.value == "openai":
            try:
                from langchain_openai import ChatOpenAI

                self.model_instance = ChatOpenAI(model=self.model.name, **model_config)
            except ImportError as e:
                raise ImportError(
                    "OpenAI provider requires 'langchain-openai'. Install it with: pip install langchain-openai"
                ) from e

        elif self.model.provider.value == "anthropic":
            try:
                from langchain_anthropic import ChatAnthropic

                self.model_instance = ChatAnthropic(model=self.model.name, **model_config)
            except ImportError as e:
                raise ImportError(
                    "Anthropic provider requires 'langchain-anthropic'. "
                    "Install it with: pip install langchain-anthropic"
                ) from e
        else:
            raise ValueError(f"Unsupported model provider: {self.model.provider}")

        # Get all available tools
        all_tools = self._inbuilt_tools() + self.agent_tools()

        # Create parsers
        output_parser = None
        tools_parser = None

        # Create PydanticOutputParser for response format instructions
        if isinstance(self.output_schema, type) and issubclass(self.output_schema, BaseModel):
            output_parser = PydanticOutputParser(pydantic_object=self.output_schema)

        # Create PydanticToolsParser for tool parsing if we have Pydantic-based tools
        if all_tools:
            from langchain_core.output_parsers import PydanticToolsParser

            # Extract Pydantic tool schemas if available
            pydantic_tools = []
            for tool_func in all_tools:
                if hasattr(tool_func, "args_schema") and tool_func.args_schema:
                    if isinstance(tool_func.args_schema, type) and issubclass(tool_func.args_schema, BaseModel):
                        pydantic_tools.append(tool_func.args_schema)

            if pydantic_tools:
                tools_parser = PydanticToolsParser(tools=pydantic_tools)

        # Build system prompt with format instructions (properly escaped)
        system_prompt = self._build_system_prompt(all_tools, output_parser)

        # Create prompt template with support for additional_instructions
        self.prompt_template = ChatPromptTemplate.from_messages(
            [
                ("system", system_prompt),
                ("system", "{additional_instructions}"),  # Dynamic instructions injected at runtime
                ("placeholder", "{messages}"),
            ]
        )

        # Build the chain based on what parsers we have
        if all_tools:
            llm_with_tools = self.model_instance.bind_tools(all_tools)
        else:
            llm_with_tools = self.model_instance

        # Create the appropriate chain based on available parsers
        if output_parser and tools_parser:
            # Both structured output and tools - need to handle this carefully
            if isinstance(self.output_schema, type) and issubclass(self.output_schema, BaseModel):
                structured_llm = llm_with_tools.with_structured_output(self.output_schema)
            else:
                structured_llm = llm_with_tools
            self.runtime = self.prompt_template | structured_llm

        elif tools_parser and not output_parser:
            # Only tools, no structured output
            self.runtime = self.prompt_template | llm_with_tools | tools_parser

        elif output_parser and not tools_parser:
            # Only structured output, no tool parser needed
            if isinstance(self.output_schema, type) and issubclass(self.output_schema, BaseModel):
                structured_llm = llm_with_tools.with_structured_output(self.output_schema)
            else:
                structured_llm = llm_with_tools
            self.runtime = self.prompt_template | structured_llm

        else:
            # No parsers needed, simple chain
            self.runtime = self.prompt_template | llm_with_tools

    def _build_system_prompt(self, available_tools: list[Callable], parser: PydanticOutputParser = None) -> str:
        """Build comprehensive system prompt with tool descriptions and format instructions."""
        prompt = self.instruction + "\n\n" + BASE_INSTRUCTIONS

        if available_tools:
            prompt += "\n\nYou have access to the following tools:\n"
            for tool_func in available_tools:
                if hasattr(tool_func, "name") and hasattr(tool_func, "description"):
                    prompt += f"- {tool_func.name}: {tool_func.description}\n"
                elif hasattr(tool_func, "__name__") and hasattr(tool_func, "__doc__"):
                    prompt += f"- {tool_func.__name__}: {tool_func.__doc__ or 'No description'}\n"

            prompt += "\nUse tools when necessary to provide accurate and helpful responses."

        # Add format instructions but escape JSON braces to avoid template variable conflicts
        if parser:
            format_instructions = parser.get_format_instructions()
            # Escape curly braces to prevent them from being treated as template variables
            format_instructions = format_instructions.replace("{", "{{").replace("}", "}}")
            prompt += f"\n\nWrap your response in `json` tags and follow this format:\n{format_instructions}"

        return prompt

    def get_identity(self):
        return {"id": self.id, "name": self.name, "label": self.label, "type": "agent"}

    async def _deliver_to_channels(self, raw_output, thread_id, message_id, run_id):
        if not self.channels or not self.channels_agent:
            return
        tasks = []
        for ch in self.channels:
            if not ch.enabled:
                continue
            tasks.append(
                self.channels_agent.run(
                    channel=ch,
                    output=raw_output,
                    thread_id=thread_id,
                    message_id=message_id,
                    run_id=run_id,
                ),
            )
        if tasks:
            await asyncio.gather(*tasks)

    async def run(
        self,
        message: str,
        user_id: str | None = None,
        thread_id: str | None = None,
        message_id: str | None = None,
        history_turns: int | None = None,  # Override agent's default history_turns
    ):
        # Ensure stable, consistent IDs with proper fallbacks
        user_id = user_id or "anonymous"
        thread_id = thread_id or str(uuid.uuid4())
        message_id = message_id or str(uuid.uuid4())
        run_id = str(uuid.uuid4())

        # Use provided history_turns or fall back to agent's configured value
        effective_history_turns = history_turns if history_turns is not None else self.history_turns

        with ls.trace(
            name=f"Agent Run: {self.label}",
            run_type="chain",
            inputs={"message": message, "user_id": user_id},
            tags=[f"{self.label}", user_id],
            metadata={
                "thread_id": thread_id,
                "user_id": user_id,
                "run_id": run_id,
                "history_turns_requested": effective_history_turns,
            },
        ) as rt:
            try:
                rt.metadata["started_at"] = str(
                    datetime.datetime.now(datetime.timezone.utc),
                )

                result = await self._run_direct_llm(
                    message, thread_id, message_id, user_id, effective_history_turns, run_id
                )

                return result

            except Exception as e:
                raise e

            finally:
                rt.metadata["completed_at"] = str(
                    datetime.datetime.now(datetime.timezone.utc),
                )

    async def _run_direct_llm(
        self, message: str, thread_id: str, message_id: str, user_id: str, effective_history_turns: int, run_id: str
    ):
        """Execute using direct LLM approach with proper message handling."""

        # Determine if this is a sub-agent (for proper message role tagging)
        is_subagent = isinstance(self, SubAgent)
        agent_type = "sub_agent" if is_subagent else "agent"

        # Build conversation context with history
        messages_context = []
        additional_instructions = ""
        history_count = 0

        # Only fetch history if agent is conversational
        if self.conversational and effective_history_turns > 0:
            try:
                history_items = await self.get_thread_history(thread_id, limit=effective_history_turns)
                history_count = len(history_items)

                # DEBUG: Log what we retrieved
                logger.info(
                    f"[Conversational Mode] Retrieved {history_count} history items for thread {thread_id[:8]}..."
                )

                # Build additional_instructions with conversation context (OpenAIAssistantRunnable pattern)
                if history_count > 0:
                    additional_instructions = (
                        f"CONVERSATION CONTEXT: You are continuing an ongoing conversation in "
                        f"thread '{thread_id[:8]}."
                        f"..'. "
                        f"Below are the {history_count} most recent message exchanges from this conversation. "
                        f"Review this history carefully and maintain consistency, continuity, and context awareness "
                        f"with your previous responses. Reference prior exchanges when relevant to provide a cohesive"
                        f" experience. "
                        f"The conversation history is ordered chronologically from oldest to newest."
                    )

                # Convert history items to LangChain message format
                for item in sorted(history_items, key=lambda x: x.value.get("_created_at", "")):
                    msg_data = item.value
                    role = msg_data.get("role", "user")
                    content = msg_data.get("content", "")

                    # DEBUG: Log each message being added
                    logger.info(f"  Adding to context: [{role}] {content[:50]}...")

                    # Skip sub_agent messages from history (they don't contribute to conversational context)
                    if role == "sub_agent":
                        continue
                    elif role == "user":
                        messages_context.append(HumanMessage(content=content))
                    elif role == "assistant":
                        messages_context.append(AIMessage(content=content))
                    elif role == "system":
                        messages_context.append(SystemMessage(content=content))

            except Exception as e:
                # If history fetch fails, continue without history
                logger.warning(f"Failed to fetch history for thread {thread_id}: {e}")
                pass
        else:
            # Non-conversational agent - skip history entirely
            logger.info(f"[{agent_type.upper()}] Skipping conversation history for thread {thread_id[:8]}...")

        # Set appropriate additional_instructions based on conversational mode
        if not additional_instructions:
            if self.conversational:
                additional_instructions = (
                    "This is the start of a new conversation. No previous message history is available."
                )
            else:
                additional_instructions = "Process this request independently without conversation history."

        # DEBUG: Log final context
        logger.info(f"Final context has {len(messages_context)} messages before adding current message")
        logger.info(f"Additional instructions: {additional_instructions[:100]}...")

        # Add current message
        messages_context.append(HumanMessage(content=message))

        # Persist incoming message with appropriate role
        # Sub-agents use role="sub_agent" for both query and response
        incoming_role = "sub_agent" if is_subagent else "user"

        await self.save_thread_message(
            thread_id,
            {
                "role": incoming_role,
                "content": message,
                "message_id": message_id,
                "user_id": user_id,
                "agent_id": self.id,  # Track which agent/sub-agent
                "agent_type": agent_type,
            },
        )

        # Build input dict following OpenAIAssistantRunnable API pattern
        input_dict = {
            "messages": messages_context,
            "additional_instructions": additional_instructions,
        }

        # Invoke the LLM chain with proper RunnableConfig
        config = {
            "recursion_limit": self.recursion_limit,
            "tags": [f"agent_{self.id}", f"user_{user_id}", f"thread_{thread_id[:8]}"],
            "metadata": {
                "agent_id": self.id,
                "agent_label": self.label,
                "user_id": user_id,
                "thread_id": thread_id,
                "message_id": message_id,
                "run_id": run_id,
                "history_turns_effective": effective_history_turns,
                "history_count": history_count,
                "conversational": self.conversational,
                "agent_type": agent_type,
            },
            "configurable": {
                "thread_id": thread_id,
                "user_id": user_id,
            },
        }

        result = await self.runtime.ainvoke(input_dict, config=config)

        # Smart content extraction for conversation history
        content = self._extract_readable_content(result)

        # Persist outgoing response with appropriate role
        outgoing_role = "sub_agent" if is_subagent else "assistant"

        await self.save_thread_message(
            thread_id,
            {
                "role": outgoing_role,
                "content": content,  # Clean, human-readable content for LLM context
                "message_id": str(uuid.uuid4()),
                "user_id": user_id,
                "agent_id": self.id,  # Track which agent/sub-agent
                "agent_type": agent_type,
                "structured_response": result.model_dump() if hasattr(result, "model_dump") else result,
                # Full structured response preserved separately
            },
        )

        # Deliver to channels if enabled
        if self.deliver_to_channels:
            asyncio.create_task(
                self._deliver_to_channels(
                    result,
                    thread_id,
                    message_id,
                    str(uuid.uuid4()),
                ),
            )

        return result

    def _extract_readable_content(self, result: Any) -> str:
        """
        Extract human-readable content from structured outputs for conversation history.
        This ensures conversation history is clean and understandable for the LLM.

        Strategy:
        1. For BaseModel with 'result' field -> extract result
        2. For BaseModel with descriptive fields -> create summary
        3. For dict with 'result' -> extract result
        4. Otherwise -> string representation
        """
        if isinstance(result, BaseModel):
            # Try to extract 'result' field if it exists
            if hasattr(result, "result"):
                return str(result.result)

            # For complex structured outputs, create a readable summary
            # Look for common descriptive fields
            model_dict = result.model_dump() if hasattr(result, "model_dump") else {}

            # Common patterns for structured outputs
            summary_fields = []

            # Check for common descriptive fields in order of preference
            priority_fields = [
                "message_to_recipient",
                "message",
                "summary",
                "description",
                "justification",
                "reason",
                "explanation",
                "response",
                "classification",
                "status",
                "title",
            ]

            for field in priority_fields:
                if field in model_dict and model_dict[field]:
                    summary_fields.append(f"{field}: {model_dict[field]}")

            # If we found descriptive fields, use them
            if summary_fields:
                return " | ".join(summary_fields[:3])  # Limit to top 3 fields

            # Otherwise, create a compact representation
            # Filter out None/empty values and limit fields
            compact_dict = {k: v for k, v in model_dict.items() if v is not None and v != ""}
            if len(compact_dict) <= 3:
                return str(compact_dict)
            else:
                # Show first 3 fields + indicator of more
                items = list(compact_dict.items())[:3]
                content = ", ".join([f"{k}={v}" for k, v in items])
                return f"{content} ... (+{len(compact_dict) - 3} more fields)"

        else:
            # For dict responses, try to extract the 'result' field
            if isinstance(result, dict):
                if "result" in result:
                    return str(result["result"])
                # For other dicts, create readable string
                return str(result)
            else:
                {"result": str(result)}
                return str(result)

    @abstractmethod
    def agent_tools(self) -> list[Callable]:
        """Subclass must return list of tools to include."""
        raise NotImplementedError


class SubAgent(BaseAgent):
    def __init__(
        self,
        config: SubAgentSchema,
        parent: Agent = None,
        deliver_to_channels: bool = False,
    ):
        super().__init__(
            name=config.name,
            model=config.model or (parent.model if parent else None),
            tools=config.tools,
            strict_mode=config.strict_mode,
            label=config.description or config.name,
            instruction=config.prompt or "You are a specialized sub-agent.",
            output_schema=config.output_schema,
            deliver_to_channels=deliver_to_channels,
            history_turns=getattr(config, "history_turns", 0),  # Use config value, default to 0
            conversational=getattr(config, "conversational", False),  # Use config value, default to False
        )
        self.config = config
        self.parent = parent
        self.initialize_runtime()
        self.channels: list[ChannelSchema] = config.channels
        self.channels_agent = ChannelsAgent(self.model) if self.channels else None

    def agent_tools(self) -> list[Callable]:
        return self.tools

    async def run(
        self,
        message: str,
        user_id: str | None = None,
        thread_id: str | None = None,
        message_id: str | None = None,
        history_turns: int | None = None,
    ):
        """
        Override run() to use role="sub_agent" for messages.
        Sub-agents save to parent's thread for audit trail.
        History access is controlled by config.history_turns.
        """
        # Inherit parent's thread_id if available
        if thread_id is None and self.parent:
            # Get thread_id from parent's current context (if running)
            # For now, we'll need to pass it explicitly from the routing tool
            thread_id = str(uuid.uuid4())

        # Use provided history_turns or respect agent's configured value
        # Don't force to 0 - let the config control this
        if history_turns is None:
            history_turns = self.history_turns

        return await super().run(
            message=message,
            user_id=user_id,
            thread_id=thread_id,
            message_id=message_id,
            history_turns=history_turns,
        )


class Agent(BaseAgent):
    def __init__(self, config: AgentSchema, deliver_to_channels: bool = True):
        super().__init__(
            name=config.name,
            model=config.model,
            tools=config.tools,
            strict_mode=config.strict_mode,
            label=config.description or config.name,
            instruction=config.instruction,
            output_schema=config.output_schema,
            deliver_to_channels=deliver_to_channels,
            history_turns=getattr(config, "history_turns", 10),
            conversational=getattr(config, "conversational", True),
        )
        self.config = config
        self.channels: list[ChannelSchema] = config.channels
        self.channels_agent = ChannelsAgent(self.model) if self.channels else None

        # Build sub-agent router
        self.subagent_router = {
            sa.name: SubAgent(sa, parent=self, deliver_to_channels=False) for sa in config.subagents
        }

        # Create single routing tool
        if self.subagent_router:
            self.subagent_routing_tool = self._build_routing_tool()
        else:
            self.subagent_routing_tool = None

        self.initialize_runtime()

    def _build_routing_tool(self) -> Callable:
        """
        Create a single delegation tool for routing to specialized sub-agents.
        Sub-agents run in the same thread with configurable history access.
        """
        from langchain_core.tools import tool

        # Build structured description with sub-agent metadata
        specialists_desc_parts = []
        for sa_id, sa in self.subagent_router.items():
            meta_parts = [
                f"  • {sa_id}",
                f"    Description: {sa.config.description}",
            ]
            if sa.config.capabilities:
                meta_parts.append(f"    Capabilities: {', '.join(sa.config.capabilities)}")
            if sa.config.tags:
                meta_parts.append(f"    Tags: {', '.join(sa.config.tags)}")
            specialists_desc_parts.append("\n".join(meta_parts))

        specialists_desc = "\n\n".join(specialists_desc_parts)

        # Capture parent agent reference for context access
        parent_agent = self

        @tool(
            description=(
                f"Delegate complex queries to specialized sub-agents when their expertise is needed.\n\n"
                f"Available specialists:\n{specialists_desc}\n\n"
                f"Use this when the query requires specialized analysis beyond your core capabilities."
            ),
            parse_docstring=True,
        )
        async def delegate_to_specialist(
            specialist_id: str,
            query: str,
            thread_id: str = None,
        ) -> dict[str, Any]:
            """
            Delegate query to a specialized sub-agent.

            Args:
                specialist_id: ID of the specialist sub-agent (exact name from available specialists)
                query: The specific query or task to send to the specialist
                thread_id: Thread ID for message persistence (auto-captured from parent context)

            Returns:
                dict: The specialist's structured analysis with capability metadata
            """
            if specialist_id not in parent_agent.subagent_router:
                available = ", ".join(parent_agent.subagent_router.keys())
                logger.warning(f"Unknown specialist '{specialist_id}'. Available: {available}")
                return {
                    "status": "error",
                    "error": f"Unknown specialist '{specialist_id}'. Available: {available}",
                    "specialist_id": specialist_id,
                }

            sub_agent = parent_agent.subagent_router[specialist_id]

            # Extract metadata for logging and return
            specialist_capabilities = sub_agent.config.capabilities

            try:
                logger.info(f"Delegating to specialist '{specialist_id}' with history_turns={sub_agent.history_turns}")

                # Sub-agent runs with parent's thread for message continuity
                # Uses subagent's configured history_turns
                result = await sub_agent.run(
                    message=query,
                    user_id=None,
                    thread_id=thread_id,
                    message_id=str(uuid.uuid4()),
                    history_turns=sub_agent.history_turns,
                )

                logger.info(f"Specialist '{specialist_id}' completed successfully")

                return {
                    "status": "success",
                    "specialist_id": specialist_id,
                    "specialist_capabilities": specialist_capabilities,
                    "result": result,
                }

            except Exception as e:
                logger.error(f"Specialist '{specialist_id}' failed: {e}", exc_info=True)
                return {
                    "status": "error",
                    "specialist_id": specialist_id,
                    "error": str(e),
                    "message": f"Specialist encountered an error during analysis: {str(e)}",
                }

        return delegate_to_specialist

    def agent_tools(self) -> list[Callable]:
        from .tools.memory import build_memory_tools

        tools = self.tools.copy()

        # Add routing tool if sub-agents exist
        if self.subagent_routing_tool:
            tools.append(self.subagent_routing_tool)

        # Add memory tools
        tools.extend(build_memory_tools(self))

        return tools
