#!/usr/bin/env python3
"""
Text-to-GraphQL Agent Implementation using LangGraph.
This file contains the implementation of the agent that processes natural language
queries into GraphQL queries, validates them, executes them, and visualizes the results.
"""

from typing import Dict, Any, List, Optional, TypedDict, Literal, cast, Union
import json
import re

from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.language_models import BaseChatModel

from .config import settings
from .logger import logger
from .types import AgentState

# Import tools
from .tools import (
    intent_recognition,
    schema_management,
    intent_options,
    select_options,
    query_construction,
    query_validation,
    query_execution,
    data_visualization,
    query_optimization
)

# Initialize LLM
llm = ChatOpenAI(
    model=settings.MODEL_NAME,
    temperature=settings.MODEL_TEMPERATURE,
    api_key=settings.OPENAI_API_KEY
)

# Tool definitions
tools = [
    {
        "name": "intent_recognition",
        "description": "Analyzes natural language queries to identify GraphQL intent (query vs mutation)",
        "when_to_use": "Always use as the first step of processing a new user query to determine what the user is trying to accomplish",
        "prerequisites": "Natural language query",
        "provides": "Intent information (query/mutation)"
    },
    {
        "name": "schema_management",
        "description": "Manages GraphQL schema and identifies relevant schema types for the user's query",
        "when_to_use": "After intent recognition to identify the relevant parts of the GraphQL schema",
        "prerequisites": "Intent information",
        "provides": "Relevant schema types and their fields"
    },
    {
        "name": "intent_options",
        "description": "Gets the available query options based on the user's intent and schema",
        "when_to_use": "After schema management to generate a list of possible GraphQL operations",
        "prerequisites": "Intent information and schema",
        "provides": "List of possible operations"
    },
    {
        "name": "select_options",
        "description": "Selects the most appropriate GraphQL operation from the available options",
        "when_to_use": "After intent_options to select the best option for the user's query",
        "prerequisites": "List of options from intent_options",
        "provides": "Selected GraphQL operation"
    },
    {
        "name": "query_construction",
        "description": "Builds executable GraphQL queries from intent and selected option",
        "when_to_use": "After selecting an option to construct the actual GraphQL query",
        "prerequisites": "Selected GraphQL operation",
        "provides": "GraphQL query string"
    },
    {
        "name": "query_validation",
        "description": "Validates GraphQL queries for correctness and optimization",
        "when_to_use": "After query construction to ensure the query is valid",
        "prerequisites": "GraphQL query string",
        "provides": "Validation result and possibly an improved query"
    },
    {
        "name": "query_optimization",
        "description": "Optimizes GraphQL queries based on validation errors",
        "when_to_use": "After query validation if the query is invalid and needs improvement",
        "prerequisites": "GraphQL query and validation errors",
        "provides": "Improved GraphQL query"
    },
    {
        "name": "query_execution",
        "description": "Executes GraphQL queries against the API endpoint",
        "when_to_use": "After query validation if the query is valid",
        "prerequisites": "Valid GraphQL query",
        "provides": "Query execution results"
    },
    {
        "name": "data_visualization",
        "description": "Recommends visualizations for query results",
        "when_to_use": "After query execution if there is data to visualize",
        "prerequisites": "Query execution results",
        "provides": "Visualization recommendations and configurations"
    }
]

# Function implementations
tool_functions = {
    "intent_recognition": intent_recognition,
    "schema_management": schema_management,
    "intent_options": intent_options,
    "select_options": select_options,
    "query_construction": query_construction,
    "query_validation": query_validation,
    "query_optimization": query_optimization,
    "query_execution": query_execution,
    "data_visualization": data_visualization
}

# LLM tool selector
def get_tool_selector_prompt() -> ChatPromptTemplate:
    """Returns the prompt for the LLM tool selector"""
    template = """You are the controller for a GraphQL agent. Your task is to determine which tool the agent should use next based on the current state of its execution.

## Available Tools
{tools}

## Current Agent State
{state}

## Guidelines
- Analyze the current state to determine what information the agent has and what it needs next
- Consider prerequisites for each tool and whether they're satisfied
- Select the most appropriate tool for the current state
- If the workflow is complete or has encountered an error that can't be resolved, select "END"
- Provide a confidence score between 0 and 1 indicating your certainty in this decision
  - 0.0-0.3 for low confidence (unsure which tool to use)
  - 0.4-0.7 for medium confidence (reasonable guess)
  - 0.8-1.0 for high confidence (clear choice)

## Response Format
Respond with a JSON object with three fields:
1. "tool": The name of the tool to use next (one of: "intent_recognition", "schema_management", "intent_options", "select_options", "query_construction", "query_validation", "query_execution", "data_visualization", or "END")
2. "reasoning": A brief explanation of why this tool is appropriate at this stage
3. "confidence": A number between 0 and 1 indicating your confidence in this selection

Example response:
```json
{
  "tool": "schema_management",
  "reasoning": "Intent recognition has been completed, so now we need to identify the relevant schema types.",
  "confidence": 0.9
}
```
"""
    return ChatPromptTemplate.from_template(template)

async def tool_selector(state: AgentState) -> Dict[str, Any]:
    """Decides which tool to use next based on the current state"""
    try:
        # Get the LLM from state
        agent_llm = state.get("default_llm")
        if not agent_llm:
            logger.warning("No default LLM provided for tool selection")
            return {"tool": "END", "reasoning": "No LLM available for tool selection", "confidence": 0.0}
            
        # Format the state for the prompt
        # Remove large data structures and LLM object that might cause prompt size issues
        state_for_prompt = {k: v for k, v in state.items() if k not in ["default_llm", "messages"]}
        
        # If there's execution_result data, summarize it instead of including all of it
        if "execution_result" in state_for_prompt:
            if isinstance(state_for_prompt["execution_result"], dict):
                keys = list(state_for_prompt["execution_result"].get("data", {}).keys())
                state_for_prompt["execution_result"] = {
                    "status": state_for_prompt["execution_result"].get("status", "UNKNOWN"),
                    "data_keys": keys,
                    "has_data": bool(state_for_prompt["execution_result"].get("data"))
                }
        
        # Similarly, summarize visualization results
        if "visualization_result" in state_for_prompt:
            viz_result = state_for_prompt["visualization_result"]
            if isinstance(viz_result, dict):
                viz_count = len(viz_result.get("visualizations", []))
                viz_types = [v.get("type") for v in viz_result.get("visualizations", [])]
                state_for_prompt["visualization_result"] = {
                    "visualization_count": viz_count,
                    "visualization_types": viz_types,
                    "has_recommendations": bool(viz_result.get("recommendations"))
                }
        
        # Get current step info
        current_step = state.get("current_step", "start")
        
        # Function to make rule-based decisions with high confidence
        def rule_based_decision() -> Dict[str, Any]:
            # Simple rule-based decisions based on workflow state to handle common transitions
            if current_step == "intent_recognition" and state.get("intent"):
                return {"tool": "schema_management", "reasoning": "Intent recognized, moving to schema management", "confidence": 0.95}
                
            if current_step == "schema_management" and state.get("schema"):
                return {"tool": "intent_options", "reasoning": "Schema retrieved, getting intent options", "confidence": 0.95}
                
            if current_step == "intent_options" and state.get("options"):
                return {"tool": "select_options", "reasoning": "Options retrieved, selecting the best option", "confidence": 0.95}
                
            if current_step == "select_options" and state.get("selected_option"):
                return {"tool": "query_construction", "reasoning": "Option selected, constructing query", "confidence": 0.95}
                
            if current_step == "query_construction" and state.get("query"):
                return {"tool": "query_validation", "reasoning": "Query constructed, validating", "confidence": 0.95}
                
            if current_step == "query_validation":
                validation_result = state.get("query_validation_result", {}).get("is_valid", False)
                if validation_result and state.get("query"):
                    # return {"tool": "query_execution", "reasoning": "Query validated, executing", "confidence": 0.95}
                    # Going to let the execute_query endpoint handle this instead of doing it automatically here
                    return {"tool": "END", "reasoning": "Query validated, returning results", "confidence": 0.95}

                else:
                    # Check if we should try optimization
                    optimization_attempts = state.get("optimization_attempts", 0)
                    max_optimization_attempts = 3
                    if optimization_attempts < max_optimization_attempts and validation_result is False:
                        return {"tool": "query_optimization", "reasoning": "Query validation failed, attempting optimization", "confidence": 0.95}
                    else:
                        return {"tool": "END", "reasoning": "Query validation failed and max optimization attempts reached", "confidence": 0.9}
            
            if current_step == "query_optimization" and state.get("query"):
                return {"tool": "query_validation", "reasoning": "Query optimized, validating again", "confidence": 0.95}
                
            if current_step == "query_execution" and state.get("execution_result", {}).get("data"):
                return {"tool": "data_visualization", "reasoning": "Query executed, visualizing results", "confidence": 0.95}
                
            if current_step == "data_visualization" and state.get("visualization_result"):
                return {"tool": "END", "reasoning": "Visualization complete, ending workflow", "confidence": 0.95}
                
            # If no rule matches, return with low confidence to indicate LLM should make decision
            return None
        
        # First, try to make a rule-based decision for standard transitions
        rule_decision = rule_based_decision()
        
        # If we have a high-confidence rule-based decision, use it without consulting the LLM
        if rule_decision:
            return rule_decision
            
        # If we don't have a clear rule, consult the LLM
        # Use a more concise version of tools for the prompt
        tools_for_prompt = []
        for tool in tools:
            tools_for_prompt.append(
                f"- {tool['name']}: {tool['description']}\n  When to use: {tool['when_to_use']}\n  Prerequisites: {tool['prerequisites']}\n"
            )
        
        # Get the tool selector prompt
        prompt = get_tool_selector_prompt().format(
            tools="\n".join(tools_for_prompt),
            state=json.dumps(state_for_prompt, indent=2, default=str)
        )
        
        # Call the LLM to select the next tool
        response = await agent_llm.ainvoke(prompt)
        response_content = response.content
        
        # Parse the response
        logger.info(f"Tool selector response: {response_content}")
        
        # Extract JSON from the response
        try:
            # First try: Look for JSON code blocks
            json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', response_content)
            if json_match:
                json_str = json_match.group(1).strip()
                result = json.loads(json_str)
            else:
                # Second try: Try to parse the whole response
                result = json.loads(response_content.strip())
                
            # Ensure confidence value exists and is a float between 0 and 1
            if "confidence" not in result:
                result["confidence"] = 0.5  # Default medium confidence
            else:
                try:
                    confidence = float(result["confidence"])
                    # Ensure confidence is in valid range
                    result["confidence"] = max(0.0, min(1.0, confidence))
                except:
                    result["confidence"] = 0.5  # Default if conversion fails
            
            logger.info(f"Selected tool: {result.get('tool')} - Confidence: {result.get('confidence')} - Reasoning: {result.get('reasoning')}")
            
            # Only use LLM judgment if confidence is at least 0.4
            if result.get("confidence", 0) >= 0.4:
                return result
            else:
                logger.info(f"Low LLM confidence ({result.get('confidence')}), falling back to rule-based decision")
                # Try rule-based decision again
                fallback = rule_based_decision()
                if fallback:
                    return fallback
                
                # If no rule applies, use the LLM decision anyway, but with a note
                result["reasoning"] = f"Using LLM decision despite low confidence: {result.get('reasoning')}"
                return result
                
        except json.JSONDecodeError:
            # If we can't parse JSON, look for tool name in the response text
            logger.warning("Failed to parse JSON from LLM response, extracting tool name from text")
            
            # Try to extract confidence value from text
            confidence_match = re.search(r'["\']confidence["\']\s*:\s*(0\.\d+|1\.0|1)', response_content)
            confidence = 0.3  # Default low confidence when parsing fails
            if confidence_match:
                try:
                    confidence = float(confidence_match.group(1))
                except:
                    pass
            
            # Check for exact tool names in the response
            for tool_name in tool_functions.keys():
                if f'"tool": "{tool_name}"' in response_content or f"'tool': '{tool_name}'" in response_content:
                    return {"tool": tool_name, "reasoning": "Extracted from response text", "confidence": confidence}
            
            # Look for END keyword
            if '"tool": "END"' in response_content or "'tool': 'END'" in response_content:
                return {"tool": "END", "reasoning": "Extracted END from response text", "confidence": confidence}
                
            # Since we have parsing issues and low confidence, fall back to rule-based
            fallback = rule_based_decision()
            if fallback:
                return fallback
                
            # Try to guess the next logical step based on state
            if state.get("intent") and not state.get("schema"):
                return {"tool": "schema_management", "reasoning": "Logical next step after intent recognition", "confidence": 0.7}
            elif state.get("schema") and not state.get("options"):
                return {"tool": "intent_options", "reasoning": "Logical next step after schema management", "confidence": 0.7}
            elif state.get("options") and not state.get("selected_option"):
                return {"tool": "select_options", "reasoning": "Logical next step after getting options", "confidence": 0.7}
            elif state.get("selected_option") and not state.get("query"):
                return {"tool": "query_construction", "reasoning": "Logical next step after selecting option", "confidence": 0.7}
            elif state.get("query") and not state.get("query_validation_result"):
                return {"tool": "query_validation", "reasoning": "Logical next step after query construction", "confidence": 0.7}
            # elif state.get("query_validation_result", {}).get("is_valid") and not state.get("execution_result"):
            #     return {"tool": "query_execution", "reasoning": "Logical next step after validation", "confidence": 0.7}
            elif state.get("execution_result") and not state.get("visualization_result"):
                return {"tool": "data_visualization", "reasoning": "Logical next step after execution", "confidence": 0.7}
            else:
                return {"tool": "END", "reasoning": "Cannot determine next step, ending workflow", "confidence": 0.5}
                
    except Exception as e:
        logger.error(f"Error in tool selector: {str(e)}", exc_info=True)
        # Make a best guess at the next tool based on current state
        current_step = state.get("current_step", "start")
        
        # Simple fallback logic for common transitions
        if current_step == "intent_recognition":
            return {"tool": "schema_management", "reasoning": "Fallback to schema management after intent recognition", "confidence": 0.6}
        elif current_step == "schema_management":
            return {"tool": "intent_options", "reasoning": "Fallback to intent options after schema management", "confidence": 0.6}
        elif current_step == "intent_options":
            return {"tool": "select_options", "reasoning": "Fallback to select options after intent options", "confidence": 0.6}
        elif current_step == "select_options":
            return {"tool": "query_construction", "reasoning": "Fallback to query construction after select options", "confidence": 0.6}
        elif current_step == "query_construction":
            return {"tool": "query_validation", "reasoning": "Fallback to query validation after construction", "confidence": 0.6}
        else:
            return {"tool": "END", "reasoning": f"Error in tool selector: {str(e)}. Ending workflow.", "confidence": 0.5}

# Agent controller (decides which tool to use next)
async def agent_controller(state: AgentState) -> Literal[
    "intent_recognition", "schema_management", "intent_options", 
    "select_options", "query_construction", "query_validation", 
    "query_optimization", "query_execution", "data_visualization", END
]:
    """Determines the next step in the agent's workflow using the tool selector"""
    
    # Check for tool errors
    if state.get("tool_error"):
        logger.error(f"Tool error encountered: {state['tool_error']}")
        return END
    
    # If this is the start, always begin with intent recognition
    current_step = state.get("current_step", "start")
    if current_step == "start":
        return "intent_recognition"
    
    # Otherwise, use the tool selector to determine the next step
    selector_result = await tool_selector(state)
    next_tool = selector_result.get("tool")
    confidence = selector_result.get("confidence", 0.0)
    
    # Log the selection with confidence
    logger.json("INFO", "Tool selection", {
        "current_step": current_step,
        "selected_tool": next_tool,
        "confidence": confidence,
        "reasoning": selector_result.get("reasoning")
    })
    
    # Return the selected tool or END
    if next_tool == "END":
        return END
    elif next_tool in tool_functions:
        return cast(
            Literal["intent_recognition", "schema_management", "intent_options", 
                   "select_options", "query_construction", "query_validation", 
                   "query_optimization", "query_execution", "data_visualization"],
            next_tool
        )
    else:
        logger.warning(f"Unknown tool selected: {next_tool}, ending workflow")
        return END

# Create a graph for the agent workflow
def build_graph() -> StateGraph:
    """Build the workflow graph for the GraphQL agent"""
    
    workflow = StateGraph(AgentState)
    
    # Add nodes for each step in the workflow
    for tool_name, tool_function in tool_functions.items():
        workflow.add_node(tool_name, tool_function)
    
    # Add a specific start node to resolve the empty string issue
    workflow.add_node("start", lambda x: x)
    
    # Add conditional edges using the async controller
    for node_name in list(tool_functions.keys()) + ["start"]:
        workflow.add_conditional_edges(node_name, agent_controller)
    
    # Set entry point to the start node
    workflow.set_entry_point("start")
    
    logger.json("INFO", "Flexible agent workflow graph built", {
        "nodes": ["start"] + list(tool_functions.keys()),
        "entry_point": "start",
        "recursion_limit": settings.RECURSION_LIMIT
    })
    
    return workflow

# Main GraphQL Agent class
class GraphQLAgent:
    """LangGraph-based GraphQL Agent with flexible workflow"""
    
    def __init__(self):
        """Initialize the GraphQL Agent"""
        self.graph = build_graph().compile()
        logger.info("Flexible GraphQL Agent initialized")
        self.llm = ChatOpenAI(
            model=settings.MODEL_NAME,
            temperature=settings.MODEL_TEMPERATURE,
            api_key=settings.OPENAI_API_KEY
        )
        
        logger.json("INFO", "LLM configuration", {
            "model": settings.MODEL_NAME,
            "temperature": settings.MODEL_TEMPERATURE,
            "tracing_enabled": False
        })

        self.state = {
            "default_llm": self.llm,
            "current_step": "start",
            "messages": []
        }
    
    async def generate_query(self, natural_language_query: str, session_id: str) -> Dict[str, Any]:
        """Generate a GraphQL query from natural language"""
        logger.llm("Processing natural language query", prompt=natural_language_query)

        # Workaround due to new queries automatically pulling up previous query result instead of generating a new one
        # Probably need a better solution in the future
        self.state = {
            "default_llm": self.llm,
            "current_step": "start",
            "messages": []
        }

        # Initialize the state for query generation flow
        self.state['messages'].append(HumanMessage(content=natural_language_query))
        self.state['natural_language_query'] = natural_language_query
        # Reset optimization attempts for a new query
        self.state['optimization_attempts'] = 0
        # Execute the graph
        result = await self.graph.ainvoke(self.state)
        self.state = result
        
        # Extract relevant information
        response = {
            "graphql_query": result.get("query"),
            "validation_result": result.get("query_validation_result")
        }
        
        logger.json("INFO", "Query generation completed", response)
        return response

    async def validate_query(self, graphql_query: str, natural_language_query: str, session_id) -> Dict[str, Any]:
        """Validate a GraphQL query"""
        logger.json("INFO", "Validating GraphQL query", {
            "graphql_query": graphql_query[:100] + ("..." if len(graphql_query) > 100 else ""),
            "natural_language_query": natural_language_query
        })
        
        self.state['messages'].append(HumanMessage(content=f"Validate this GraphQL query: {graphql_query}"))
        self.state['natural_language_query'] = natural_language_query
        self.state['query'] = graphql_query
        
        # Execute just the validation step
        result = await query_validation(self.state)
        self.state.update(result)
        
        # Return validation result
        validation_result = result.get("query_validation_result", {})
        logger.json("INFO", "Query validation result", validation_result)
        return validation_result

    async def execute_query(self, graphql_query: str, variables: Dict[str, Any] = None) -> Dict[str, Any]:
        """Execute a GraphQL query"""
        logger.json("INFO", "Executing GraphQL query", {
            "graphql_query": graphql_query[:100] + ("..." if len(graphql_query) > 100 else "")
        })
        
        self.state['messages'].append(HumanMessage(content=f"Execute this GraphQL query: {graphql_query}"))
        self.state['query'] = graphql_query

        # Execute just the execution step
        if variables:
            result = await query_execution(self.state, variables=variables)
        else:
            result = await query_execution(self.state)
        self.state.update(result)
        
        # Return execution result
        execution_result = result.get("execution_result", {})
        return execution_result

    async def visualize_data(self, data: Dict[str, Any], natural_language_query: str) -> Dict[str, Any]:
        """Generate visualizations for query results"""
        logger.json("INFO", "Generating visualizations", {
            "data_size": len(str(data)),
            "natural_language_query": natural_language_query,
            "data_structure": list(data.keys()) if isinstance(data, dict) else "non-dict data"
        })
        
        self.state['messages'].append(HumanMessage(content=f"Visualize this data: {json.dumps(data)}"))
        self.state['execution_result'] = {"data": data, "status": "SUCCESS"}
        self.state['current_step'] = "data_visualization"

        # Execute just the visualization step
        result = await data_visualization(self.state)
        self.state.update(result)
        
        # Return visualization result
        viz_result = result.get("visualization_result", {})
        logger.json("INFO", "Visualization generated", {
            "num_visualizations": len(viz_result.get("visualizations", [])),
            "visualization_types": [v.get("type") for v in viz_result.get("visualizations", [])],
            "recommendations": viz_result.get("recommendations", []),
            "visualizations": viz_result.get("visualizations", [])
        })
        return viz_result