"""
VSS LangGraph Orchestrator
Mode 1 query workflow as a state machine.

This file shows the complete VSSOrchestrator class with the updated run() method
that properly handles media input.
"""

import logging
from typing import Literal, Optional, Dict, Any
from datetime import datetime, timezone

from langgraph.graph import StateGraph, END

from orchestration.state import VSSState, create_state, add_trace, set_error
from agents.routing_agent import classify_intent
from agents.sql_agent import process_sql
from agents.vlm_agent import process_vlm
from agents.response_composer import compose_response
from config.settings import get_settings
from config.logging_config import log_node_step, log_request_trace

logger = logging.getLogger(__name__)


class VSSOrchestrator:
    """
    LangGraph orchestrator for Mode 1 (User Query) workflow.
    
    Flow:
        classify_intent
             │
        ┌────┴────┐
        ▼         ▼
      SQL        VLM
        │         │
        ├─escalate┤
        ▼         ▼
        compose_response
             │
            END
    """
    
    def __init__(self):
        self.settings = get_settings()
        self.graph = self._build_graph()
        self.app = self.graph.compile()
    
    def _build_graph(self) -> StateGraph:
        """Build the LangGraph state machine."""
        graph = StateGraph(VSSState)
        
        # Nodes
        graph.add_node("classify", self._classify)
        graph.add_node("sql", self._sql)
        graph.add_node("vlm", self._vlm)
        graph.add_node("compose", self._compose)
        graph.add_node("error", self._error)
        
        # Entry
        graph.set_entry_point("classify")
        
        # Routing from classify
        graph.add_conditional_edges(
            "classify",
            self._route_intent,
            {"sql": "sql", "vlm": "vlm", "error": "error"}
        )
        
        # Routing from SQL
        graph.add_conditional_edges(
            "sql",
            self._route_sql,
            {"compose": "compose", "escalate": "vlm", "error": "error"}
        )
        
        # Routing from VLM
        graph.add_conditional_edges(
            "vlm",
            self._route_vlm,
            {"compose": "compose", "error": "error"}
        )
        
        # Terminal edges
        graph.add_edge("compose", END)
        graph.add_edge("error", END)
        
        return graph
    
    # === Node Wrappers ===
    
    async def _classify(self, state: VSSState) -> VSSState:
        """Intent classification node."""
        prev = len(state.get("trace", []))
        add_trace(state, "Node: classify")
        state["status"] = "routing"
        result = await classify_intent(state)
        log_node_step(result, "classify", prev)
        return result

    async def _sql(self, state: VSSState) -> VSSState:
        """SQL processing node."""
        prev = len(state.get("trace", []))
        add_trace(state, "Node: sql")
        state["response_source"] = "sql"
        result = await process_sql(state)
        log_node_step(result, "sql", prev)
        return result

    async def _vlm(self, state: VSSState) -> VSSState:
        """VLM processing node."""
        prev = len(state.get("trace", []))
        add_trace(state, "Node: vlm")
        state["response_source"] = "vlm"
        result = await process_vlm(state)
        log_node_step(result, "vlm", prev)
        return result

    async def _compose(self, state: VSSState) -> VSSState:
        """Response composition node."""
        prev = len(state.get("trace", []))
        add_trace(state, "Node: compose")
        state = await compose_response(state)
        state["status"] = "completed"

        # Calculate processing time
        try:
            start = datetime.fromisoformat(state["timestamp"])
            now = datetime.now(timezone.utc)
            # Ensure both datetimes are timezone-aware for subtraction
            if start.tzinfo is None:
                start = start.replace(tzinfo=timezone.utc)
            state["processing_time_ms"] = int((now - start).total_seconds() * 1000)
        except Exception:
            pass

        log_node_step(state, "compose", prev)
        return state

    async def _error(self, state: VSSState) -> VSSState:
        """Error handling node."""
        prev = len(state.get("trace", []))
        add_trace(state, "Node: error")
        error = state.get("error", "Unknown error")
        state["final_response"] = f"Error processing request: {error}"
        state["status"] = "failed"
        state["response_source"] = "error"
        log_node_step(state, "error", prev)
        return state
    
    # === Routing Functions ===
    
    def _route_intent(self, state: VSSState) -> Literal["sql", "vlm", "error"]:
        """Route based on classified intent. VLM path is only taken when vlm_enabled=True."""
        if state.get("error"):
            return "error"
        intent = state.get("intent", "unknown")
        if intent == "vlm":
            if self.settings.vlm_enabled:
                return "vlm"
            # VLM disabled — fall through to SQL and note the redirect
            add_trace(state, "Routing: VLM intent detected but VLM disabled — redirecting to SQL")
            state["routing_reason"] = "VLM disabled; redirected to SQL"
        return "sql"
    
    def _route_sql(self, state: VSSState) -> Literal["compose", "escalate", "error"]:
        """Route after SQL processing. Escalation to VLM requires both vlm_enabled and enable_vlm_escalation."""
        if state.get("error"):
            return "error"
        can_escalate = self.settings.vlm_enabled and self.settings.enable_vlm_escalation
        if state.get("escalated") and can_escalate:
            return "escalate"
        if state.get("sql_error") and can_escalate:
            state["escalated"] = True
            state["escalation_reason"] = state["sql_error"]
            return "escalate"
        return "compose"
    
    def _route_vlm(self, state: VSSState) -> Literal["compose", "error"]:
        """Route after VLM processing."""
        if state.get("error") or state.get("vlm_error"):
            if not state.get("error"):
                set_error(state, state.get("vlm_error", "VLM failed"))
            return "error"
        return "compose"
    
    # === Public API ===
    
    async def run(
        self,
        query: str,
        media: Optional[Dict[str, Any]] = None,
        camera_id: Optional[str] = None,
        options: Optional[Dict[str, Any]] = None,
    ) -> VSSState:
        """
        Execute query through orchestration workflow.
        
        Args:
            query: Natural language query string
            media: Optional media dict with keys:
                   - type: "image" or "video"
                   - data: base64 encoded content
                   - filename: optional original filename
            camera_id: Optional camera filter
            options: Optional query options
        
        Returns:
            Final VSSState with response and metadata
        """
        # Create initial state
        state = create_state(query)
        state["camera_id"] = camera_id
        state["options"] = options or {}
        
        # Add media to state if provided
        if media:
            state["uploaded_media"] = media
            add_trace(state, f"Media attached: {media.get('type', 'unknown')}")
        
        logger.info(f"Starting orchestration: {state['request_id']}")
        add_trace(state, f"Query: {query[:50]}...")
        
        try:
            # Run the graph
            result = await self.app.ainvoke(state)
            logger.info(f"Orchestration complete: {result['request_id']} - {result['status']}")
            log_request_trace(result)
            return result
            
        except Exception as e:
            logger.error(f"Orchestration failed: {e}", exc_info=True)
            state["error"] = str(e)
            state["status"] = "failed"
            state["final_response"] = f"System error: {str(e)}"
            add_trace(state, f"FATAL: {e}")
            log_request_trace(state)
            return state


# === Module-level convenience function ===

_orchestrator: Optional[VSSOrchestrator] = None


async def run_query(
    query: str,
    media: Optional[Dict[str, Any]] = None,
) -> VSSState:
    """
    Run a query through the VSS orchestrator.

    Convenience function that manages the orchestrator singleton.
    """
    global _orchestrator
    if _orchestrator is None:
        _orchestrator = VSSOrchestrator()
    return await _orchestrator.run(query=query, media=media)



# """
# VSS LangGraph Orchestrator
# Mode 1 query workflow as a state machine.
# """

# import logging
# from typing import Literal, Optional
# from datetime import datetime

# from langgraph.graph import StateGraph, END

# from orchestration.state import VSSState, create_state, add_trace, set_error
# from agents.routing_agent import classify_intent
# from agents.sql_agent import process_sql
# from agents.vlm_agent import process_vlm
# from agents.response_composer import compose_response
# from config.settings import get_settings

# logger = logging.getLogger(__name__)


# class VSSOrchestrator:
#     """
#     LangGraph orchestrator for Mode 1 (User Query) workflow.
    
#     Flow:
#         classify_intent
#              │
#         ┌────┴────┐
#         ▼         ▼
#       SQL        VLM
#         │         │
#         ├─escalate┤
#         ▼         ▼
#         compose_response
#              │
#             END
#     """
    
#     def __init__(self):
#         self.settings = get_settings()
#         self.graph = self._build_graph()
#         self.app = self.graph.compile()
    
#     def _build_graph(self) -> StateGraph:
#         graph = StateGraph(VSSState)
        
#         # Nodes
#         graph.add_node("classify", self._classify)
#         graph.add_node("sql", self._sql)
#         graph.add_node("vlm", self._vlm)
#         graph.add_node("compose", self._compose)
#         graph.add_node("error", self._error)
        
#         # Entry
#         graph.set_entry_point("classify")
        
#         # Routing from classify
#         graph.add_conditional_edges(
#             "classify",
#             self._route_intent,
#             {"sql": "sql", "vlm": "vlm", "error": "error"}
#         )
        
#         # Routing from SQL
#         graph.add_conditional_edges(
#             "sql",
#             self._route_sql,
#             {"compose": "compose", "escalate": "vlm", "error": "error"}
#         )
        
#         # Routing from VLM
#         graph.add_conditional_edges(
#             "vlm",
#             self._route_vlm,
#             {"compose": "compose", "error": "error"}
#         )
        
#         # Terminal edges
#         graph.add_edge("compose", END)
#         graph.add_edge("error", END)
        
#         return graph
    
#     # === Node Wrappers ===
    
#     async def _classify(self, state: VSSState) -> VSSState:
#         add_trace(state, "Node: classify")
#         state["status"] = "routing"
#         return await classify_intent(state)
    
#     async def _sql(self, state: VSSState) -> VSSState:
#         add_trace(state, "Node: sql")
#         state["response_source"] = "sql"
#         return await process_sql(state)
    
#     async def _vlm(self, state: VSSState) -> VSSState:
#         add_trace(state, "Node: vlm")
#         state["response_source"] = "vlm"
#         return await process_vlm(state)
    
#     async def _compose(self, state: VSSState) -> VSSState:
#         add_trace(state, "Node: compose")
#         state = await compose_response(state)
#         state["status"] = "completed"
        
#         # Calculate processing time
#         try:
#             start = datetime.fromisoformat(state["timestamp"])
#             state["processing_time_ms"] = int((datetime.utcnow() - start).total_seconds() * 1000)
#         except:
#             pass
        
#         return state
    
#     async def _error(self, state: VSSState) -> VSSState:
#         add_trace(state, "Node: error")
#         error = state.get("error", "Unknown error")
#         state["final_response"] = f"Error processing request: {error}"
#         state["status"] = "failed"
#         state["response_source"] = "error"
#         return state
    
#     # === Routing Functions ===
    
#     def _route_intent(self, state: VSSState) -> Literal["sql", "vlm", "error"]:
#         if state.get("error"):
#             return "error"
#         intent = state.get("intent", "unknown")
#         if intent == "vlm":
#             return "vlm"
#         return "sql"  # Default to SQL
    
#     def _route_sql(self, state: VSSState) -> Literal["compose", "escalate", "error"]:
#         if state.get("error"):
#             return "error"
#         if state.get("escalated") and self.settings.enable_vlm_escalation:
#             return "escalate"
#         if state.get("sql_error") and self.settings.enable_vlm_escalation:
#             state["escalated"] = True
#             state["escalation_reason"] = state["sql_error"]
#             return "escalate"
#         return "compose"
    
#     def _route_vlm(self, state: VSSState) -> Literal["compose", "error"]:
#         if state.get("error") or state.get("vlm_error"):
#             if not state.get("error"):
#                 set_error(state, state.get("vlm_error", "VLM failed"))
#             return "error"
#         return "compose"
    
#     # === Public API ===
    
#     async def run(self, query: str, media: Optional[dict] = None) -> VSSState:
#         """
#         Execute query through orchestration.
        
#         Args:
#             query: Natural language question
#             media: Optional {"type": "image", "data": bytes}
        
#         Returns:
#             Final state with response
#         """
#         state = create_state(query, media)
#         logger.info(f"Starting query: {state['request_id']}")
        
#         try:
#             result = await self.app.ainvoke(state)
#             logger.info(f"Completed: {result['request_id']} status={result['status']}")
#             return result
#         except Exception as e:
#             logger.error(f"Orchestration failed: {e}")
#             set_error(state, str(e))
#             state["final_response"] = f"Unexpected error: {e}"
#             return state


# # Singleton
# _orchestrator: Optional[VSSOrchestrator] = None

# def get_orchestrator() -> VSSOrchestrator:
#     global _orchestrator
#     if _orchestrator is None:
#         _orchestrator = VSSOrchestrator()
#     return _orchestrator


# async def run_query(query: str, media: Optional[dict] = None) -> VSSState:
#     """Convenience function to run a query."""
#     return await get_orchestrator().run(query, media)