"""
Advanced Examples & Patterns for Multi-Agent Procurement System

Demonstrates:
1. Custom conditional routing
2. Error handling and recovery
3. Async execution
4. State persistence queries
5. Parallel vendor evaluation
6. Budget refinement loops
"""

import asyncio
from typing import Optional
from datetime import datetime
import json

from langchain_groq import ChatGroq
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import Command

from procurement_system import (
    ProcurementState,
    ProcurementAgents,
    build_procurement_graph,
    ProcurementWorkflowExecutor,
    create_llm,
)


# =============================================================================
# EXAMPLE 1: Multi-Stage Approval with Escalation
# =============================================================================

class EscalationApprovalGate:
    """
    Enhanced approval gate with escalation logic:
    - Budget < $5000: Auto-approve
    - Budget $5000-$10000: Manager approval
    - Budget > $10000: Executive approval
    """

    @staticmethod
    def determine_approval_level(budget: float) -> str:
        """
        Route approval based on budget threshold.

        Args:
            budget: Procurement budget

        Returns:
            Approval level: "auto", "manager", or "executive"
        """
        if budget < 5000:
            return "auto"
        elif budget <= 10000:
            return "manager"
        else:
            return "executive"

    @staticmethod
    def approval_gate_escalation(state: ProcurementState) -> Command:
        """
        Advanced approval gate with multi-level routing.

        Args:
            state: Current procurement state

        Returns:
            Command routing to appropriate approval level
        """
        budget = state["budget_limit"]
        approval_level = EscalationApprovalGate.determine_approval_level(budget)

        new_log = f"[{datetime.now().isoformat()}] EscalationGate: Budget ${budget:,.2f} requires {approval_level.upper()} approval"

        # Route based on approval level
        if approval_level == "auto":
            # Auto-approve small procurements
            return Command(
                update={
                    "human_approved": True,
                    "logs": [new_log + " (AUTO-APPROVED)"]
                },
                goto="legal_node"
            )
        elif approval_level == "manager":
            return Command(
                update={"logs": [new_log]},
                goto="manager_approval"
            )
        else:  # executive
            return Command(
                update={"logs": [new_log]},
                goto="executive_approval"
            )


# =============================================================================
# EXAMPLE 2: Budget Refinement Loop
# =============================================================================

class BudgetRefinement:
    """
    Handles budget overages by suggesting refinements and re-routing to analysis.
    """

    @staticmethod
    def refinement_node(state: ProcurementState, llm: ChatGroq) -> Command:
        """
        When budget is exceeded, suggest alternatives to procurement requester.

        Refinement options:
        1. Reduce scope (fewer features/users)
        2. Select lower-cost vendor
        3. Negotiate terms with selected vendor
        4. Increase budget

        Args:
            state: Current state with budget-exceeded vendor
            llm: Language model for generating suggestions

        Returns:
            Command with refinement options
        """
        vendor = state["selected_vendor"]
        budget = state["budget_limit"]
        shortfall = vendor.get("price_per_month", 0) - budget

        new_log = f"[{datetime.now().isoformat()}] RefinementNode: Budget shortfall ${shortfall:,.2f}"

        # Use LLM to generate refinement options
        from langchain_core.prompts import PromptTemplate

        refinement_prompt = PromptTemplate(
            input_variables=["vendor", "budget", "shortfall"],
            template="""
Budget refinement required.

Selected vendor: {vendor_name} - ${vendor_price}/month
Budget limit: ${budget}
Shortfall: ${shortfall}

Generate 3 refinement options for the procurement requester:
1. Scope reduction suggestions
2. Alternative vendor from list
3. Negotiation talking points for vendor

Format as JSON with "options" array.
"""
        )

        prompt_text = refinement_prompt.format(
            vendor_name=vendor.get("name", "Unknown"),
            vendor_price=vendor.get("price_per_month", 0),
            budget=budget,
            shortfall=shortfall
        )

        response = llm.invoke(prompt_text)
        refinement_options = response.content

        new_log += " | Refinement options generated"

        # In production, these would be sent to procurement requester
        # They would choose an option, state would update, and graph would continue

        return Command(
            update={
                "logs": [new_log],
                "contract_draft": f"Refinement Options:\n{refinement_options}"
            },
            goto=END
        )


# =============================================================================
# EXAMPLE 3: Async Workflow Execution
# =============================================================================

async def execute_workflow_async(
    graph,
    procurement_request: str,
    budget_limit: float
) -> dict:
    """
    Async execution of procurement workflow for concurrent processing.

    Enables:
    - Running multiple workflows in parallel
    - Non-blocking I/O for approval waits
    - Better resource utilization

    Args:
        graph: Compiled LangGraph
        procurement_request: Procurement need
        budget_limit: Budget limit

    Returns:
        Final workflow state
    """
    thread_id = f"async_procurement_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    config = {"configurable": {"thread_id": thread_id}}

    initial_state = {
        "procurement_request": procurement_request,
        "vendor_options": [],
        "selected_vendor": {},
        "budget_limit": budget_limit,
        "analysis_approved": False,
        "human_approved": False,
        "contract_draft": "",
        "logs": ["Async workflow initiated"]
    }

    # Stream execution asynchronously
    final_state = None
    async for event in graph.astream(initial_state, config):
        print(f"Async Event: {event}")

    # Get state at interruption
    final_state = graph.get_state(config)
    return final_state.values


async def run_multiple_workflows():
    """
    Execute multiple procurement workflows concurrently.

    Use case: Process multiple department requests in parallel.
    """
    llm = create_llm()
    graph, _, _ = build_procurement_graph(llm)

    # Create multiple concurrent workflows
    workflows = [
        execute_workflow_async(graph, "Cloud infrastructure", 5000),
        execute_workflow_async(graph, "Software licensing", 3000),
        execute_workflow_async(graph, "Security tools", 2000),
    ]

    # Execute all in parallel
    results = await asyncio.gather(*workflows)

    print("\n" + "="*80)
    print("CONCURRENT WORKFLOWS COMPLETED")
    print("="*80)
    for i, result in enumerate(results, 1):
        print(f"Workflow {i}: {result['selected_vendor'].get('name', 'N/A')}")
    print("="*80)

    return results


# =============================================================================
# EXAMPLE 4: State Query & History Tracking
# =============================================================================

class WorkflowStateInspector:
    """
    Query and analyze workflow state at any point in execution.
    """

    def __init__(self, graph, thread_id: str):
        """
        Initialize inspector for a specific workflow.

        Args:
            graph: Compiled LangGraph
            thread_id: Thread ID of workflow to inspect
        """
        self.graph = graph
        self.thread_id = thread_id
        self.config = {"configurable": {"thread_id": thread_id}}

    def get_current_state(self) -> dict:
        """Retrieve current state of workflow."""
        state = self.graph.get_state(self.config)
        return state.values

    def get_state_history(self) -> list:
        """Retrieve all state snapshots for this workflow."""
        # Note: MemorySaver doesn't track history by default
        # Production with PostgresSaver would query checkpoint history
        current = self.get_current_state()
        return [current]

    def get_decision_trail(self) -> list:
        """Extract decision points from logs."""
        state = self.get_current_state()
        logs = state.get("logs", [])

        decisions = []
        for log in logs:
            if "Route" in log or "APPROVED" in log or "REJECTED" in log:
                decisions.append(log)

        return decisions

    def estimate_contract_ready_date(self) -> Optional[str]:
        """
        Based on current state, estimate when contract will be ready.
        """
        state = self.get_current_state()

        if state.get("contract_draft"):
            return "READY"  # Contract already generated
        elif state.get("human_approved"):
            return "PENDING"  # Awaiting legal node execution
        elif state.get("analysis_approved"):
            return "AWAITING_APPROVAL"  # Waiting for human
        else:
            return "BLOCKED"  # Budget or analysis issue

    def generate_audit_report(self) -> str:
        """Generate comprehensive audit report of workflow."""
        state = self.get_current_state()

        report = f"""
PROCUREMENT WORKFLOW AUDIT REPORT
Generated: {datetime.now().isoformat()}
Thread ID: {self.thread_id}

PROCUREMENT DETAILS
├─ Request: {state.get('procurement_request')[:100]}...
├─ Budget: ${state.get('budget_limit', 0):,.2f}
└─ Status: {"COMPLETE" if state.get('contract_draft') else "IN_PROGRESS"}

DECISION HISTORY
"""
        for i, log in enumerate(state.get("logs", []), 1):
            report += f"├─ {i}. {log}\n"

        report += f"""
VENDOR SELECTION
├─ Candidates Evaluated: {len(state.get('vendor_options', []))}
├─ Selected: {state.get('selected_vendor', {}).get('name', 'N/A')}
├─ Price: ${state.get('selected_vendor', {}).get('price_per_month', 0):,.2f}/month
└─ Status: {"APPROVED" if state.get('human_approved') else "PENDING_APPROVAL"}

CONTRACT STATUS
└─ Generated: {"YES" if state.get('contract_draft') else "NO"}

END REPORT
"""
        return report


# =============================================================================
# EXAMPLE 5: Vendor Comparison Matrix
# =============================================================================

class VendorComparison:
    """
    Generate detailed vendor comparison for procurement stakeholders.
    """

    @staticmethod
    def build_comparison_matrix(state: ProcurementState, llm: ChatGroq) -> str:
        """
        Create side-by-side vendor comparison.

        Args:
            state: State with vendor_options populated
            llm: Language model for analysis

        Returns:
            Markdown formatted comparison table
        """
        vendors = state.get("vendor_options", [])
        budget = state.get("budget_limit", 0)

        from langchain_core.prompts import PromptTemplate

        comparison_prompt = PromptTemplate(
            input_variables=["vendors", "budget"],
            template="""
Create a markdown comparison matrix for these vendors against budget ${budget}.

Vendors:
{vendors}

Generate a table comparing:
- Vendor name
- Price per month
- Reputation score
- Key capabilities
- Within budget (✓/✗)
- Risk assessment

Format as markdown table.
"""
        )

        vendors_json = json.dumps(vendors, indent=2)
        prompt_text = comparison_prompt.format(vendors=vendors_json, budget=budget)

        response = llm.invoke(prompt_text)
        return response.content


# =============================================================================
# EXAMPLE 6: Error Recovery with Retry Logic
# =============================================================================

def robust_research_node(state: ProcurementState, llm: ChatGroq, max_retries: int = 3) -> Command:
    """
    Research node with built-in retry logic and error handling.

    Args:
        state: Current state
        llm: Language model
        max_retries: Maximum retry attempts

    Returns:
        Command with state updates
    """
    from procurement_system import mock_vendor_search

    request = state["procurement_request"]
    new_logs = []

    for attempt in range(max_retries):
        try:
            vendors = mock_vendor_search(request)
            if not vendors:
                raise ValueError("No vendors returned from search")

            new_log = f"[{datetime.now().isoformat()}] ResearchNode: Attempt {attempt + 1}: Found {len(vendors)} vendors"
            new_logs.append(new_log)

            return Command(
                update={
                    "vendor_options": vendors,
                    "logs": new_logs
                },
                goto="analysis_node"
            )

        except Exception as e:
            error_log = f"[{datetime.now().isoformat()}] ResearchNode: Attempt {attempt + 1} failed: {str(e)}"
            new_logs.append(error_log)

            if attempt == max_retries - 1:
                # Final attempt failed
                new_logs.append("ResearchNode: All retry attempts exhausted")
                return Command(
                    update={"logs": new_logs},
                    goto=END
                )
            # Retry


# =============================================================================
# EXAMPLE 7: Custom Metrics & KPIs
# =============================================================================

class ProcurementMetrics:
    """
    Track and analyze procurement workflow metrics.
    """

    @staticmethod
    def calculate_procurement_cycle_time(state: ProcurementState) -> float:
        """
        Calculate time from request to approval.

        Args:
            state: Final state with logs

        Returns:
            Cycle time in seconds
        """
        logs = state.get("logs", [])
        if len(logs) < 2:
            return 0

        # Extract timestamps from logs
        try:
            first_time = logs[0].split("]")[0].strip("[")
            last_time = logs[-1].split("]")[0].strip("[")

            from datetime import datetime
            t1 = datetime.fromisoformat(first_time)
            t2 = datetime.fromisoformat(last_time)
            return (t2 - t1).total_seconds()
        except:
            return 0

    @staticmethod
    def calculate_vendor_variance(state: ProcurementState) -> dict:
        """
        Calculate variance between vendor prices.

        Args:
            state: State with vendor_options

        Returns:
            Variance metrics
        """
        vendors = state.get("vendor_options", [])
        if not vendors:
            return {}

        prices = [v.get("price_per_month", 0) for v in vendors]
        avg_price = sum(prices) / len(prices)
        min_price = min(prices)
        max_price = max(prices)

        return {
            "avg_price": avg_price,
            "min_price": min_price,
            "max_price": max_price,
            "range": max_price - min_price,
            "variance_percentage": ((max_price - min_price) / avg_price * 100) if avg_price > 0 else 0
        }

    @staticmethod
    def calculate_budget_utilization(state: ProcurementState) -> float:
        """
        Calculate percentage of budget utilized by selected vendor.

        Args:
            state: Final state

        Returns:
            Utilization percentage
        """
        vendor = state.get("selected_vendor", {})
        budget = state.get("budget_limit", 1)

        vendor_price = vendor.get("price_per_month", 0)
        utilization = (vendor_price / budget) * 100

        return round(utilization, 2)


# =============================================================================
# EXAMPLE 8: Integration with External Systems
# =============================================================================

class ExternalSystemIntegration:
    """
    Integrate procurement system with external platforms.
    """

    @staticmethod
    async def send_approval_notification_via_slack(
        state: ProcurementState,
        slack_webhook: str
    ) -> bool:
        """
        Send approval request via Slack when graph pauses.

        Args:
            state: Current procurement state
            slack_webhook: Slack webhook URL

        Returns:
            Success status
        """
        import requests

        vendor = state.get("selected_vendor", {})
        message = {
            "text": "🛒 Procurement Approval Required",
            "blocks": [
                {
                    "type": "section",
                    "text": {
                        "type": "mrkdwn",
                        "text": f"""*Vendor:* {vendor.get('name')}
*Price:* ${vendor.get('price_per_month'):,.2f}/month
*Reputation:* {vendor.get('reputation_score')}/10
*Approve?*"""
                    }
                },
                {
                    "type": "actions",
                    "elements": [
                        {
                            "type": "button",
                            "text": {"type": "plain_text", "text": "Approve"},
                            "value": "approve",
                            "style": "primary"
                        },
                        {
                            "type": "button",
                            "text": {"type": "plain_text", "text": "Reject"},
                            "value": "reject",
                            "style": "danger"
                        }
                    ]
                }
            ]
        }

        try:
            response = requests.post(slack_webhook, json=message)
            return response.status_code == 200
        except Exception as e:
            print(f"Slack notification failed: {str(e)}")
            return False

    @staticmethod
    async def sync_contract_to_docusign(
        contract_draft: str,
        recipient_email: str
    ) -> bool:
        """
        Sync generated contract to DocuSign for e-signature.

        Args:
            contract_draft: Generated contract markdown
            recipient_email: Email of contract recipient

        Returns:
            Success status
        """
        # Mock implementation
        print(f"[DocuSign] Uploading contract for {recipient_email}")
        # In production, would call DocuSign API
        return True


# =============================================================================
# MAIN: Run Examples
# =============================================================================

async def main():
    """Run advanced examples."""

    print("\n" + "="*80)
    print("ADVANCED PROCUREMENT SYSTEM EXAMPLES")
    print("="*80 + "\n")

    # Initialize LLM and graph
    llm = create_llm()
    graph, memory, agents = build_procurement_graph(llm)

    # Example 1: Basic workflow
    print("Example 1: Basic Workflow Execution\n")
    executor = ProcurementWorkflowExecutor(graph, memory)
    state, _ = executor.start_workflow(
        procurement_request="Cloud infrastructure with auto-scaling",
        budget_limit=5500.0,
    )

    # Example 2: State inspection
    print("\nExample 2: Workflow State Inspection\n")
    inspector = WorkflowStateInspector(graph, executor.thread_id)
    print("Decision Trail:")
    for decision in inspector.get_decision_trail():
        print(f"  - {decision}")

    print(f"\nEstimated Contract Ready: {inspector.estimate_contract_ready_date()}")

    # Example 3: Metrics
    print("\nExample 3: Procurement Metrics\n")
    metrics = ProcurementMetrics()
    print(f"Budget Utilization: {metrics.calculate_budget_utilization(state)}%")
    variance = metrics.calculate_vendor_variance(state)
    if variance:
        print(f"Vendor Price Variance: {variance['variance_percentage']:.2f}%")

    # Example 4: Audit report
    print("\nExample 4: Audit Report\n")
    audit = inspector.generate_audit_report()
    print(audit)

    # Example 5: Async execution (if needed)
    print("\nExample 5: Note - Async execution available via run_multiple_workflows()")
    print("(Uncomment asyncio.run(run_multiple_workflows()) to test)")

    print("\n" + "="*80)
    print("ADVANCED EXAMPLES COMPLETED")
    print("="*80 + "\n")


if __name__ == "__main__":
    # Run async examples if needed
    # asyncio.run(run_multiple_workflows())

    # Run examples
    import sys
    if sys.version_info >= (3, 7):
        asyncio.run(main())
    else:
        loop = asyncio.get_event_loop()
        loop.run_until_complete(main())