"""
Base classes for persona-aware agents using LangGraph best practices.
"""

import logging
from abc import ABC, abstractmethod
from collections.abc import Sequence
from datetime import datetime
from typing import Annotated, Any

from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain_core.tools import BaseTool
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import END, START, StateGraph, add_messages
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field
from typing_extensions import TypedDict

from maverick_mcp.config.settings import get_settings
from maverick_mcp.memory.checkpointer import get_persistent_checkpointer

logger = logging.getLogger(__name__)
settings = get_settings()


class InvestorPersona(BaseModel):
    """Defines an investor persona with risk parameters."""

    name: str
    risk_tolerance: tuple[int, int] = Field(
        description="Risk tolerance range (min, max) on 0-100 scale"
    )
    position_size_max: float = Field(
        description="Maximum position size as percentage of portfolio"
    )
    stop_loss_multiplier: float = Field(
        description="Multiplier for stop loss calculation"
    )
    preferred_timeframe: str = Field(
        default="swing", description="Preferred trading timeframe: day, swing, position"
    )
    characteristics: list[str] = Field(
        default_factory=list, description="Key behavioral characteristics"
    )


# Predefined investor personas
INVESTOR_PERSONAS = {
    "conservative": InvestorPersona(
        name="Conservative",
        risk_tolerance=(
            settings.financial.risk_tolerance_conservative_min,
            settings.financial.risk_tolerance_conservative_max,
        ),
        position_size_max=settings.financial.max_position_size_conservative,
        stop_loss_multiplier=settings.financial.stop_loss_multiplier_conservative,
        preferred_timeframe="position",
        characteristics=[
            "Prioritizes capital preservation",
            "Focuses on dividend stocks",
            "Prefers established companies",
            "Long-term oriented",
        ],
    ),
    "moderate": InvestorPersona(
        name="Moderate",
        risk_tolerance=(
            settings.financial.risk_tolerance_moderate_min,
            settings.financial.risk_tolerance_moderate_max,
        ),
        position_size_max=settings.financial.max_position_size_moderate,
        stop_loss_multiplier=settings.financial.stop_loss_multiplier_moderate,
        preferred_timeframe="swing",
        characteristics=[
            "Balanced risk/reward approach",
            "Mix of growth and value",
            "Diversified portfolio",
            "Medium-term focus",
        ],
    ),
    "aggressive": InvestorPersona(
        name="Aggressive",
        risk_tolerance=(
            settings.financial.risk_tolerance_aggressive_min,
            settings.financial.risk_tolerance_aggressive_max,
        ),
        position_size_max=settings.financial.max_position_size_aggressive,
        stop_loss_multiplier=settings.financial.stop_loss_multiplier_aggressive,
        preferred_timeframe="day",
        characteristics=[
            "High risk tolerance",
            "Growth-focused",
            "Momentum trading",
            "Short-term opportunities",
        ],
    ),
    "day_trader": InvestorPersona(
        name="Day Trader",
        risk_tolerance=(
            settings.financial.risk_tolerance_day_trader_min,
            settings.financial.risk_tolerance_day_trader_max,
        ),
        position_size_max=settings.financial.max_position_size_day_trader,
        stop_loss_multiplier=settings.financial.stop_loss_multiplier_day_trader,
        preferred_timeframe="day",
        characteristics=[
            "Intraday positions only",
            "High-frequency trading",
            "Technical analysis focused",
            "Tight risk controls",
        ],
    ),
}


class BaseAgentState(TypedDict):
    """Base state for all persona-aware agents."""

    messages: Annotated[Sequence[BaseMessage], add_messages]
    persona: str
    session_id: str
    prior_findings: str  # Summary of findings from previously executed agents


class PersonaAwareTool(BaseTool):
    """Base class for tools that adapt to investor personas."""

    persona: InvestorPersona | None = None
    # State tracking
    last_analysis_time: dict[str, datetime] = {}
    analyzed_stocks: dict[str, dict] = {}
    key_price_levels: dict[str, dict] = {}
    # Cache settings
    cache_ttl: int = settings.agent.agent_cache_ttl_seconds

    def set_persona(self, persona: InvestorPersona) -> None:
        """Set the active investor persona."""
        self.persona = persona

    def adjust_for_risk(self, value: float, parameter_type: str) -> float:
        """Adjust a value based on the persona's risk profile."""
        if not self.persona:
            return value

        # Get average risk tolerance
        risk_avg = sum(self.persona.risk_tolerance) / 2
        risk_factor = risk_avg / 50  # Normalize to 1.0 at moderate risk

        # Adjust based on parameter type
        if parameter_type == "position_size":
            # Kelly Criterion-inspired sizing
            kelly_fraction = self._calculate_kelly_fraction(risk_factor)
            adjusted = value * kelly_fraction
            return min(adjusted, self.persona.position_size_max)
        elif parameter_type == "stop_loss":
            # ATR-based dynamic stops
            return value * self.persona.stop_loss_multiplier
        elif parameter_type == "profit_target":
            # Risk-adjusted targets
            return value * (2 - risk_factor)  # Conservative = lower targets
        elif parameter_type == "volatility_filter":
            # Volatility tolerance
            return value * (2 - risk_factor)  # Conservative = lower vol tolerance
        elif parameter_type == "time_horizon":
            # Holding period in days
            if self.persona.preferred_timeframe == "day":
                return 1
            elif self.persona.preferred_timeframe == "swing":
                return 5 * risk_factor  # 2.5-7.5 days
            else:  # position
                return 20 * risk_factor  # 10-30 days
        else:
            return value

    def _calculate_kelly_fraction(self, risk_factor: float) -> float:
        """Calculate position size using Kelly Criterion."""
        # Simplified Kelly: f = (p*b - q) / b
        # where p = win probability, b = win/loss ratio, q = loss probability
        # Using risk factor to adjust expected win rate
        win_probability = 0.45 + (0.1 * risk_factor)  # 45-55% base win rate
        win_loss_ratio = 2.0  # 2:1 reward/risk
        loss_probability = 1 - win_probability

        kelly = (win_probability * win_loss_ratio - loss_probability) / win_loss_ratio

        # Apply safety factor (never use full Kelly)
        safety_factor = 0.25  # Use 25% of Kelly
        return max(0, kelly * safety_factor)

    def update_analysis_data(self, symbol: str, analysis_data: dict[str, Any]):
        """Update stored analysis data for a symbol."""
        symbol = symbol.upper()
        self.analyzed_stocks[symbol] = analysis_data
        self.last_analysis_time[symbol] = datetime.now()
        if "price_levels" in analysis_data:
            self.key_price_levels[symbol] = analysis_data["price_levels"]

    def get_stock_context(self, symbol: str) -> dict[str, Any]:
        """Get stored context for a symbol."""
        symbol = symbol.upper()
        return {
            "analysis": self.analyzed_stocks.get(symbol, {}),
            "last_analysis": self.last_analysis_time.get(symbol),
            "price_levels": self.key_price_levels.get(symbol, {}),
            "cache_expired": self._is_cache_expired(symbol),
        }

    def _is_cache_expired(self, symbol: str) -> bool:
        """Check if cached data has expired."""
        last_time = self.last_analysis_time.get(symbol.upper())
        if not last_time:
            return True

        age_seconds = (datetime.now() - last_time).total_seconds()
        return age_seconds > self.cache_ttl

    def _adjust_risk_parameters(self, params: dict) -> dict:
        """Adjust parameters based on risk profile."""
        if not self.persona:
            return params

        risk_factor = sum(self.persona.risk_tolerance) / 100  # 0.1-0.9 scale

        # Apply risk adjustments based on parameter names
        adjusted = {}
        for key, value in params.items():
            if isinstance(value, int | float):
                key_lower = key.lower()
                if any(term in key_lower for term in ["stop", "support", "risk"]):
                    # Wider stops/support for conservative, tighter for aggressive
                    adjusted[key] = value * (2 - risk_factor)
                elif any(
                    term in key_lower for term in ["resistance", "target", "profit"]
                ):
                    # Lower targets for conservative, higher for aggressive
                    adjusted[key] = value * risk_factor
                elif any(term in key_lower for term in ["size", "amount", "shares"]):
                    # Smaller positions for conservative, larger for aggressive
                    adjusted[key] = self.adjust_for_risk(value, "position_size")
                elif any(term in key_lower for term in ["volume", "liquidity"]):
                    # Higher liquidity requirements for conservative
                    adjusted[key] = value * (2 - risk_factor)
                elif any(term in key_lower for term in ["volatility", "atr", "std"]):
                    # Lower volatility tolerance for conservative
                    adjusted[key] = self.adjust_for_risk(value, "volatility_filter")
                else:
                    adjusted[key] = value
            else:
                adjusted[key] = value

        return adjusted

    def _validate_risk_levels(self, data: dict) -> bool:
        """Validate if the data meets the persona's risk criteria."""
        if not self.persona:
            return True

        min_risk, max_risk = self.persona.risk_tolerance

        # Extract risk metrics
        volatility = data.get("volatility", 0)
        beta = data.get("beta", 1.0)

        # Convert to risk score (0-100)
        volatility_score = min(100, volatility * 2)  # Assume 50% vol = 100 risk
        beta_score = abs(beta - 1) * 100  # Distance from market

        # Combined risk score
        risk_score = (volatility_score + beta_score) / 2

        if risk_score < min_risk or risk_score > max_risk:
            return False

        # Persona-specific validations
        if self.persona.name == "Conservative":
            # Additional checks for conservative investors
            if data.get("debt_to_equity", 0) > 1.5:
                return False
            if data.get("current_ratio", 0) < 1.5:
                return False
            if data.get("dividend_yield", 0) < 0.02:  # Prefer dividend stocks
                return False
        elif self.persona.name == "Day Trader":
            # Day traders need high liquidity
            if data.get("average_volume", 0) < 1_000_000:
                return False
            if data.get("spread_percentage", 0) > 0.1:  # Tight spreads only
                return False

        return True

    def format_for_persona(self, data: dict) -> dict:
        """Format output data based on persona preferences."""
        if not self.persona:
            return data

        formatted = data.copy()

        # Add persona-specific insights
        formatted["persona_insights"] = {
            "suitable_for_profile": self._validate_risk_levels(data),
            "risk_adjusted_parameters": self._adjust_risk_parameters(
                data.get("parameters", {})
            ),
            "recommended_timeframe": self.persona.preferred_timeframe,
            "max_position_size": self.persona.position_size_max,
        }

        # Add risk warnings if needed
        warnings = []
        if not self._validate_risk_levels(data):
            warnings.append(f"Risk profile outside {self.persona.name} parameters")

        if data.get("volatility", 0) > 50:
            warnings.append("High volatility - consider smaller position size")

        if warnings:
            formatted["risk_warnings"] = warnings

        return formatted


class PersonaAwareAgent(ABC):
    """
    Base class for agents that adapt behavior based on investor personas.

    This follows LangGraph best practices:
    - Uses StateGraph for workflow definition
    - Implements proper node/edge patterns
    - Supports native streaming modes
    - Uses TypedDict for state management
    """

    def __init__(
        self,
        llm,
        tools: list[BaseTool],
        persona: str = "moderate",
        checkpointer: BaseCheckpointSaver | None = None,
        ttl_hours: int = 1,
    ):
        """
        Initialize a persona-aware agent.

        Args:
            llm: Language model to use
            tools: List of tools available to the agent
            persona: Investor persona name
            checkpointer: Optional checkpointer (defaults to persistent SQLite saver)
            ttl_hours: Time-to-live for memory in hours
        """
        self.llm = llm
        self.tools = tools
        self.persona = INVESTOR_PERSONAS.get(persona, INVESTOR_PERSONAS["moderate"])
        self.ttl_hours = ttl_hours

        # Set up checkpointing -- prefer persistent SQLite-backed saver
        if checkpointer is None:
            self.checkpointer = get_persistent_checkpointer()
        else:
            self.checkpointer = checkpointer

        # Configure tools with persona
        for tool in self.tools:
            if isinstance(tool, PersonaAwareTool):
                tool.set_persona(self.persona)

        # Build the graph
        self.graph = self._build_graph()

        # Track usage
        self.total_tokens = 0
        self.conversation_start = datetime.now()

    def _build_graph(self):
        """Build the LangGraph workflow."""
        # Create the graph builder
        workflow = StateGraph(self.get_state_schema())

        # Add the agent node
        workflow.add_node("agent", self._agent_node)

        # Create tool node if tools are available
        if self.tools:
            tool_node = ToolNode(self.tools)
            workflow.add_node("tools", tool_node)

            # Add conditional edge from agent
            workflow.add_conditional_edges(
                "agent",
                self._should_continue,
                {
                    # If agent returns tool calls, route to tools
                    "continue": "tools",
                    # Otherwise end
                    "end": END,
                },
            )

            # Add edge from tools back to agent
            workflow.add_edge("tools", "agent")
        else:
            # No tools, just end after agent
            workflow.add_edge("agent", END)

        # Set entry point
        workflow.add_edge(START, "agent")

        # Compile with checkpointer
        return workflow.compile(checkpointer=self.checkpointer)

    def _agent_node(self, state: dict[str, Any]) -> dict[str, Any]:
        """The main agent node that processes messages."""
        messages = state["messages"]

        # Add system message if it's the first message
        if len(messages) == 1 and isinstance(messages[0], HumanMessage):
            system_prompt = self._build_system_prompt()

            # Inject prior findings from other agents if available
            prior_findings = state.get("prior_findings", "")
            if prior_findings:
                system_prompt += (
                    f"\n\nContext from previous agent analyses:\n{prior_findings}"
                )

            messages = [SystemMessage(content=system_prompt)] + messages

        # Call the LLM
        if self.tools:
            response = self.llm.bind_tools(self.tools).invoke(messages)
        else:
            response = self.llm.invoke(messages)

        # Track tokens (simplified)
        if hasattr(response, "content"):
            self.total_tokens += len(response.content) // 4

        # Return the response
        return {"messages": [response]}

    def _should_continue(self, state: dict[str, Any]) -> str:
        """Determine whether to continue to tools or end."""
        last_message = state["messages"][-1]

        # If the LLM makes a tool call, continue to tools
        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            return "continue"

        # Otherwise we're done
        return "end"

    def _build_system_prompt(self) -> str:
        """Build system prompt based on persona."""
        base_prompt = f"""You are a financial advisor configured for a {self.persona.name} investor profile.

Risk Parameters:
- Risk Tolerance: {self.persona.risk_tolerance[0]}-{self.persona.risk_tolerance[1]}/100
- Max Position Size: {self.persona.position_size_max * 100:.1f}% of portfolio
- Stop Loss Multiplier: {self.persona.stop_loss_multiplier}x
- Preferred Timeframe: {self.persona.preferred_timeframe}

Key Characteristics:
{chr(10).join(f"- {char}" for char in self.persona.characteristics)}

Always adjust your recommendations to match this risk profile. Be explicit about risk management."""

        return base_prompt

    @abstractmethod
    def get_state_schema(self) -> type:
        """
        Get the state schema for this agent.

        Subclasses should return their specific state schema.
        """
        return BaseAgentState

    async def ainvoke(self, query: str, session_id: str, **kwargs) -> dict[str, Any]:
        """
        Invoke the agent asynchronously.

        Args:
            query: User query
            session_id: Session identifier
            **kwargs: Additional parameters

        Returns:
            Agent response
        """
        config = {
            "configurable": {"thread_id": session_id, "persona": self.persona.name}
        }

        # Merge additional config
        if "config" in kwargs:
            config.update(kwargs["config"])

        # Run the graph
        result = await self.graph.ainvoke(
            {
                "messages": [HumanMessage(content=query)],
                "persona": self.persona.name,
                "session_id": session_id,
            },
            config=config,
        )

        return self._extract_response(result)

    def invoke(self, query: str, session_id: str, **kwargs) -> dict[str, Any]:
        """
        Invoke the agent synchronously.

        Args:
            query: User query
            session_id: Session identifier
            **kwargs: Additional parameters

        Returns:
            Agent response
        """
        config = {
            "configurable": {"thread_id": session_id, "persona": self.persona.name}
        }

        # Merge additional config
        if "config" in kwargs:
            config.update(kwargs["config"])

        # Run the graph
        result = self.graph.invoke(
            {
                "messages": [HumanMessage(content=query)],
                "persona": self.persona.name,
                "session_id": session_id,
            },
            config=config,
        )

        return self._extract_response(result)

    async def astream(
        self, query: str, session_id: str, stream_mode: str = "values", **kwargs
    ):
        """
        Stream agent responses asynchronously.

        Args:
            query: User query
            session_id: Session identifier
            stream_mode: Streaming mode (values, updates, messages, custom, debug)
            **kwargs: Additional parameters

        Yields:
            Streamed chunks based on mode
        """
        config = {
            "configurable": {"thread_id": session_id, "persona": self.persona.name}
        }

        # Merge additional config
        if "config" in kwargs:
            config.update(kwargs["config"])

        # Stream the graph
        async for chunk in self.graph.astream(
            {
                "messages": [HumanMessage(content=query)],
                "persona": self.persona.name,
                "session_id": session_id,
            },
            config=config,
            stream_mode=stream_mode,
        ):
            yield chunk

    def stream(
        self, query: str, session_id: str, stream_mode: str = "values", **kwargs
    ):
        """
        Stream agent responses synchronously.

        Args:
            query: User query
            session_id: Session identifier
            stream_mode: Streaming mode (values, updates, messages, custom, debug)
            **kwargs: Additional parameters

        Yields:
            Streamed chunks based on mode
        """
        config = {
            "configurable": {"thread_id": session_id, "persona": self.persona.name}
        }

        # Merge additional config
        if "config" in kwargs:
            config.update(kwargs["config"])

        # Stream the graph
        yield from self.graph.stream(
            {
                "messages": [HumanMessage(content=query)],
                "persona": self.persona.name,
                "session_id": session_id,
            },
            config=config,
            stream_mode=stream_mode,
        )

    def _extract_response(self, result: dict[str, Any]) -> dict[str, Any]:
        """Extract the final response from graph execution."""
        messages = result.get("messages", [])

        if not messages:
            return {"content": "No response generated", "status": "error"}

        # Get the last AI message
        last_message = messages[-1]

        return {
            "content": last_message.content
            if hasattr(last_message, "content")
            else str(last_message),
            "status": "success",
            "persona": self.persona.name,
            "message_count": len(messages),
            "session_id": result.get("session_id", ""),
        }

    def get_risk_adjusted_params(
        self, base_params: dict[str, float]
    ) -> dict[str, float]:
        """Adjust parameters based on persona risk profile."""
        adjusted = {}

        for key, value in base_params.items():
            if "size" in key.lower() or "position" in key.lower():
                adjusted[key] = self.adjust_for_risk(value, "position_size")
            elif "stop" in key.lower():
                adjusted[key] = self.adjust_for_risk(value, "stop_loss")
            elif "target" in key.lower() or "profit" in key.lower():
                adjusted[key] = self.adjust_for_risk(value, "profit_target")
            else:
                adjusted[key] = value

        return adjusted

    def adjust_for_risk(self, value: float, parameter_type: str) -> float:
        """Adjust a value based on the persona's risk profile."""
        # Get average risk tolerance
        risk_avg = sum(self.persona.risk_tolerance) / 2
        risk_factor = risk_avg / 50  # Normalize to 1.0 at moderate risk

        # Adjust based on parameter type
        if parameter_type == "position_size":
            return min(value * risk_factor, self.persona.position_size_max)
        elif parameter_type == "stop_loss":
            return value * self.persona.stop_loss_multiplier
        elif parameter_type == "profit_target":
            return value * (2 - risk_factor)  # Conservative = lower targets
        else:
            return value
