"""LangGraph-based trading workflow for OpenClaw.

This module provides the TradingWorkflow class that uses LangGraph to orchestrate
multi-agent trading analysis with the following flow:

START -> [MarketAnalysis, SentimentAnalysis, FundamentalAnalysis] (parallel)
       -> BullBearDebate
       -> DecisionFusion
       -> RiskAssessment
       -> END
"""

import asyncio
from typing import Any, Dict, List, Optional

from langgraph.graph import END, START, StateGraph
from langgraph.graph.state import CompiledStateGraph
from loguru import logger

from openclaw.workflow.nodes import (
    bull_bear_debate_node,
    decision_fusion_node,
    execute_order_node,
    fundamental_analysis_node,
    market_analysis_node,
    risk_assessment_node,
    sentiment_analysis_node,
    should_continue_after_analysis,
    should_continue_after_risk,
)
from openclaw.workflow.state import TradingWorkflowState, create_initial_state


class TradingWorkflow:
    """LangGraph-based trading workflow orchestrator.

    This class manages the complete trading analysis workflow using LangGraph
    to coordinate multiple agents in a state-driven pipeline.

    Workflow Graph:
        START
          |
          v
    +-----+-----+
    |           |
    v           v
Market   Sentiment
Analysis   Analysis
    |           |
    v           v
Fundamental   |
Analysis      |
    |         |
    +----+----+
         |
         v
    BullBearDebate
         |
         v
    DecisionFusion
         |
         v
    RiskAssessment
         |
         v
        END

    Args:
        symbol: The trading symbol to analyze
        initial_capital: Initial capital for each agent (default: $1000)
        enable_parallel: Whether to run analysis nodes in parallel (default: True)
    """

    def __init__(
        self,
        symbol: str,
        initial_capital: float = 1000.0,
        enable_parallel: bool = True,
    ):
        self.symbol = symbol
        self.initial_capital = initial_capital
        self.enable_parallel = enable_parallel
        self._graph: Optional[CompiledStateGraph] = None

        logger.info(f"TradingWorkflow initialized for {symbol}")

    def _build_graph(self) -> CompiledStateGraph:
        """Build the LangGraph state graph.

        Returns:
            Compiled state graph ready for execution
        """
        # Create the state graph
        workflow = StateGraph(TradingWorkflowState)

        # Add nodes
        workflow.add_node("market_analysis", market_analysis_node)
        workflow.add_node("sentiment_analysis", sentiment_analysis_node)
        workflow.add_node("fundamental_analysis", fundamental_analysis_node)
        workflow.add_node("bull_bear_debate", bull_bear_debate_node)
        workflow.add_node("decision_fusion", decision_fusion_node)
        workflow.add_node("risk_assessment", risk_assessment_node)
        workflow.add_node("execute_order", execute_order_node)

        # Add edges from START to parallel analysis nodes
        workflow.add_edge(START, "market_analysis")
        workflow.add_edge(START, "sentiment_analysis")
        workflow.add_edge(START, "fundamental_analysis")

        # Add conditional edges from analysis nodes to debate
        # All three parallel branches converge at bull_bear_debate
        workflow.add_edge("market_analysis", "bull_bear_debate")
        workflow.add_edge("sentiment_analysis", "bull_bear_debate")
        workflow.add_edge("fundamental_analysis", "bull_bear_debate")

        # Add sequential edges: debate -> fusion -> risk -> execute_order -> END
        workflow.add_edge("bull_bear_debate", "decision_fusion")
        workflow.add_edge("decision_fusion", "risk_assessment")
        workflow.add_edge("risk_assessment", "execute_order")
        workflow.add_edge("execute_order", END)

        # Compile the graph
        compiled = workflow.compile()

        logger.info("Trading workflow graph compiled successfully")
        return compiled

    @property
    def graph(self) -> CompiledStateGraph:
        """Get the compiled graph, building if necessary."""
        if self._graph is None:
            self._graph = self._build_graph()
        return self._graph

    async def run(
        self,
        max_steps: int = 10,
        debug: bool = False,
    ) -> TradingWorkflowState:
        """Run the trading workflow.

        Args:
            max_steps: Maximum number of steps to execute
            debug: Whether to enable debug logging

        Returns:
            Final workflow state
        """
        # Create initial state
        initial_state = create_initial_state(self.symbol, self.initial_capital)

        logger.info(f"Starting trading workflow for {self.symbol}")

        try:
            # Run the graph
            final_state = await self.graph.ainvoke(
                initial_state,
                config={"recursion_limit": max_steps},
            )

            logger.info(f"Trading workflow completed for {self.symbol}")

            if debug:
                self._log_state_summary(final_state)

            return final_state

        except Exception as e:
            logger.error(f"Trading workflow failed: {e}")
            raise

    def run_sync(
        self,
        max_steps: int = 10,
        debug: bool = False,
    ) -> TradingWorkflowState:
        """Run the trading workflow synchronously.

        Args:
            max_steps: Maximum number of steps to execute
            debug: Whether to enable debug logging

        Returns:
            Final workflow state
        """
        return asyncio.run(self.run(max_steps=max_steps, debug=debug))

    def stream(
        self,
        max_steps: int = 10,
        debug: bool = False,
    ):
        """Stream the workflow execution.

        Yields state updates as they occur during workflow execution.

        Args:
            max_steps: Maximum number of steps to execute
            debug: Whether to enable debug logging

        Yields:
            State updates at each step
        """
        initial_state = create_initial_state(self.symbol, self.initial_capital)

        logger.info(f"Starting streaming workflow for {self.symbol}")

        for state_update in self.graph.stream(
            initial_state,
            config={"recursion_limit": max_steps},
        ):
            if debug:
                logger.debug(f"State update: {state_update}")
            yield state_update

    async def astream(
        self,
        max_steps: int = 10,
        debug: bool = False,
    ):
        """Async stream the workflow execution.

        Args:
            max_steps: Maximum number of steps to execute
            debug: Whether to enable debug logging

        Yields:
            State updates at each step
        """
        initial_state = create_initial_state(self.symbol, self.initial_capital)

        logger.info(f"Starting async streaming workflow for {self.symbol}")

        async for state_update in self.graph.astream(
            initial_state,
            config={"recursion_limit": max_steps},
        ):
            if debug:
                logger.debug(f"State update: {state_update}")
            yield state_update

    def get_final_decision(
        self,
        state: Optional[TradingWorkflowState] = None,
    ) -> Optional[Dict[str, Any]]:
        """Get the final trading decision from workflow state.

        Args:
            state: Workflow state (uses last run if not provided)

        Returns:
            Final decision dictionary or None
        """
        if state is None:
            # Would need to store state after run
            logger.warning("No state provided to get_final_decision")
            return None

        return state.get("final_decision")

    def _log_state_summary(self, state: TradingWorkflowState) -> None:
        """Log a summary of the workflow state.

        Args:
            state: Workflow state to summarize
        """
        from openclaw.workflow.state import get_state_summary

        summary = get_state_summary(state)

        logger.info("=" * 50)
        logger.info("WORKFLOW STATE SUMMARY")
        logger.info("=" * 50)
        logger.info(f"Symbol: {summary['symbol']}")
        logger.info(f"Current Step: {summary['current_step']}")
        logger.info(f"Completed Steps: {summary['completed_steps']}")
        logger.info(f"Reports Generated:")
        logger.info(f"  - Technical: {summary['has_technical']}")
        logger.info(f"  - Sentiment: {summary['has_sentiment']}")
        logger.info(f"  - Fundamental: {summary['has_fundamental']}")
        logger.info(f"  - Bull: {summary['has_bull']}")
        logger.info(f"  - Bear: {summary['has_bear']}")
        logger.info(f"  - Fusion: {summary['has_fusion']}")
        logger.info(f"  - Risk: {summary['has_risk']}")
        logger.info(f"  - Final Decision: {summary['has_final']}")
        logger.info(f"Errors: {summary['error_count']}")
        logger.info("=" * 50)

    def visualize(self, output_path: Optional[str] = None) -> str:
        """Generate a visualization of the workflow graph.

        Args:
            output_path: Optional path to save visualization

        Returns:
            Mermaid diagram string
        """
        # Generate mermaid diagram
        mermaid = """
```mermaid
flowchart TD
    START([START])
    END([END])

    subgraph ParallelAnalysis[Parallel Analysis]
        MA[MarketAnalysis]
        SA[SentimentAnalysis]
        FA[FundamentalAnalysis]
    end

    BB[BullBearDebate]
    DF[DecisionFusion]
    RA[RiskAssessment]

    START --> MA
    START --> SA
    START --> FA

    MA --> BB
    SA --> BB
    FA --> BB

    BB --> DF
    DF --> RA
    RA --> END
```
        """.strip()

        if output_path:
            with open(output_path, "w") as f:
                f.write(mermaid)
            logger.info(f"Workflow visualization saved to {output_path}")

        return mermaid


def run_trading_workflow(
    symbol: str,
    initial_capital: float = 1000.0,
    debug: bool = False,
) -> Dict[str, Any]:
    """Convenience function to run a complete trading workflow.

    Args:
        symbol: Trading symbol to analyze
        initial_capital: Initial capital for agents
        debug: Enable debug logging

    Returns:
        Final decision dictionary
    """
    workflow = TradingWorkflow(
        symbol=symbol,
        initial_capital=initial_capital,
    )

    final_state = workflow.run_sync(debug=debug)
    decision = workflow.get_final_decision(final_state)

    return decision or {"error": "No decision generated", "symbol": symbol}


# Export main classes and functions
__all__ = [
    "TradingWorkflow",
    "run_trading_workflow",
]
