src/visualization_workflow.py

from typing import TypedDict, List, Dict, Any, Optional, Union, Literal
import re
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from src.visualization_schema import (
    VisualizationResult, 
    EChartsConfig, 
    SeriesData, 
    AxisConfig, 
    ChartType,
    FieldInfo
)
import pandas as pd
import numpy as np
import json
import traceback
import time
import os
import io
from collections import Counter

# State definition for visualization processing workflow
class VisualizationProcessingState(TypedDict):
    # Input data
    query: str
    csv_data: str
    field_info: List[FieldInfo]
    form_title: str
    
    # Processing configuration
    processing_method: Literal["structured_llm", "fallback"]
    
    # Intermediate data
    parsed_dataframe: Optional[pd.DataFrame]
    data_summary: Optional[Dict[str, Any]]
    column_analysis: Optional[Dict[str, Any]]
    
    # Results
    visualization_result: Optional[VisualizationResult]
    error: Optional[str]
    success: bool
    processing_time: float
    
    # Metadata
    processing_complete: bool
    usage: Optional[Dict[str, int]]

class VisualizationWorkflow:
    """LangGraph workflow for data visualization generation with structured output."""
    
    def __init__(self):
        self.workflow = self._build_workflow()
    
    def _build_workflow(self) -> StateGraph:
        """Build the LangGraph workflow with conditional routing."""
        
        # Create the state graph
        workflow = StateGraph(VisualizationProcessingState)
        
        # Add nodes
        workflow.add_node("validate_input", self.validate_input_node)
        workflow.add_node("parse_data", self.parse_data_node)
        workflow.add_node("analyze_data", self.analyze_data_node)
        workflow.add_node("generate_visualization", self.generate_visualization_node)
        workflow.add_node("finalize_results", self.finalize_results_node)
        
        # Add edges
        workflow.add_edge("validate_input", "parse_data")
        workflow.add_edge("parse_data", "analyze_data")
        workflow.add_edge("analyze_data", "generate_visualization")
        workflow.add_edge("generate_visualization", "finalize_results")
        workflow.add_edge("finalize_results", END)
        
        # Set entry point
        workflow.set_entry_point("validate_input")
        
        return workflow.compile()
    
    def validate_input_node(self, state: VisualizationProcessingState) -> VisualizationProcessingState:
        """Node 1: Validate input and determine processing method."""
        try:
            # Validate required fields
            if not state.get("query") or not state.get("csv_data"):
                state["error"] = "Missing required fields: query and csv_data"
                state["success"] = False
                return state
            
            # Determine processing method
            state["processing_method"] = "structured_llm"
            
            print(f"✅ Input validated - Processing method: {state['processing_method']}")
            
        except Exception as e:
            state["error"] = f"Input validation error: {str(e)}"
            state["success"] = False
            print(f"❌ Input validation failed: {str(e)}")
        
        return state
    
    def parse_data_node(self, state: VisualizationProcessingState) -> VisualizationProcessingState:
        """Node 2: Parse CSV data into pandas DataFrame."""
        try:
            csv_data = state["csv_data"]
            
            # Handle different line endings and encoding issues
            csv_data = csv_data.replace('\r\n', '\n').replace('\r', '\n')
            df = pd.read_csv(io.StringIO(csv_data))
            
            # Clean column names (remove quotes, spaces)
            df.columns = df.columns.str.strip().str.replace('"', '')
            
            # Handle common data type conversions
            for col in df.columns:
                if df[col].dtype == 'object':
                    # Try to convert numeric columns
                    numeric_values = pd.to_numeric(df[col], errors='coerce')
                    if not numeric_values.isna().all():
                        df[col] = numeric_values
            
            state["parsed_dataframe"] = df
            
            print(f"✅ Data parsed successfully - Shape: {df.shape}, Columns: {df.columns.tolist()}")
            
        except Exception as e:
            error_msg = f"Data parsing error: {str(e)}"
            state["error"] = error_msg
            state["success"] = False
            print(f"❌ {error_msg}")
            print(traceback.format_exc())
        
        return state
    
    def analyze_data_node(self, state: VisualizationProcessingState) -> VisualizationProcessingState:
        """Node 3: Analyze the parsed data to understand structure and content."""
        try:
            df = state["parsed_dataframe"]
            if df is None or df.empty:
                state["error"] = "No data available for analysis"
                state["success"] = False
                return state
            
            # Analyze data structure
            data_summary = {
                "shape": df.shape,
                "columns": df.columns.tolist(),
                "dtypes": df.dtypes.to_dict(),
                "null_counts": df.isnull().sum().to_dict(),
                "sample_data": df.head(3).to_dict('records') if len(df) > 0 else []
            }
            
            # Analyze each column
            column_analysis = {}
            for col in df.columns:
                col_info = {
                    "dtype": str(df[col].dtype),
                    "null_count": df[col].isnull().sum(),
                    "unique_count": df[col].nunique(),
                    "sample_values": df[col].dropna().head(5).tolist()
                }
                
                # Add statistical info for numeric columns
                if pd.api.types.is_numeric_dtype(df[col]):
                    col_info.update({
                        "min": float(df[col].min()) if not df[col].isnull().all() else None,
                        "max": float(df[col].max()) if not df[col].isnull().all() else None,
                        "mean": float(df[col].mean()) if not df[col].isnull().all() else None,
                        "std": float(df[col].std()) if not df[col].isnull().all() else None
                    })
                
                column_analysis[col] = col_info
            
            state["data_summary"] = data_summary
            state["column_analysis"] = column_analysis
            
            print(f"✅ Data analysis completed - {len(df.columns)} columns analyzed")
            
        except Exception as e:
            error_msg = f"Data analysis error: {str(e)}"
            state["error"] = error_msg
            state["success"] = False
            print(f"❌ {error_msg}")
            print(traceback.format_exc())
        
        return state
    
    def generate_visualization_node(self, state: VisualizationProcessingState) -> VisualizationProcessingState:
        """Node 4: Generate visualization using structured LLM output."""
        try:
            query = state["query"]
            df = state["parsed_dataframe"]
            data_summary = state["data_summary"]
            column_analysis = state["column_analysis"]
            field_info = state["field_info"]
            form_title = state["form_title"]
            
            if df is None or df.empty:
                state["error"] = "No data available for visualization"
                state["success"] = False
                return state
            
            # Create client for code generation with structured output
            client = ChatOpenAI(
                api_key=os.getenv("OPENAI_API_KEY"), 
                model="gpt-5.4-nano", 
                temperature=0, 
                max_tokens=3000,
            )
            
            # Use structured output to ensure we get only Python code
            from pydantic import BaseModel
            
            class CodeResponse(BaseModel):
                python_code: str = "Python code that processes the DataFrame and returns result dictionary"
            
            structured_client = client.with_structured_output(CodeResponse, method="function_calling", include_raw=True)
            
            # Generate comprehensive prompt for code generation
            prompt = self._generate_code_prompt(
                query, df, data_summary, column_analysis, field_info, form_title
            )
            
            # Make API call
            start_time = time.time()
            # invoke returns a dict with 'parsed' (the model) and 'raw' (the generation output) when include_raw=True
            response = structured_client.invoke([{"role": "user", "content": prompt}])
            processing_time = time.time() - start_time
            
            # Execute the generated code
            result = self._execute_generated_code(response['parsed'].python_code, df)
            
            state["visualization_result"] = result
            state["processing_time"] = processing_time
            state["success"] = result.get('success', False)

            # Extract usage
            if "raw" in response and hasattr(response["raw"], "usage_metadata"):
                 usage = response["raw"].usage_metadata
                 state["usage"] = {
                     "input_tokens": usage.get("input_tokens", 0),
                     "output_tokens": usage.get("output_tokens", 0),
                     "total_tokens": usage.get("total_tokens", 0)
                 }
            
            print(f"✅ Visualization generation completed in {processing_time:.2f}s")
            
        except Exception as e:
            error_msg = f"Visualization generation error: {str(e)}"
            state["error"] = error_msg
            state["success"] = False
            print(f"❌ {error_msg}")
            print(traceback.format_exc())
        
        return state
    
    def _generate_code_prompt(
        self, 
        query: str, 
        df: pd.DataFrame, 
        data_summary: Dict[str, Any], 
        column_analysis: Dict[str, Any],
        field_info: List[FieldInfo],
        form_title: str
    ) -> str:
        """Generate a prompt for Python code generation that processes data and creates ECharts JSON."""
        
        # Create field type mapping
        field_types = {field.label: field.field_type for field in field_info}
        
        # Get sample data for better context
        sample_data = df.head(2).to_dict('records') if not df.empty else []
        
        prompt = f"""Generate Python code that processes the existing DataFrame 'df' and returns simple chart data.

QUERY: "{query}"
FORM: {form_title}

DATAFRAME INFO:
- Shape: {df.shape[0]} rows, {df.shape[1]} columns
- Columns: {', '.join(df.columns.tolist())}
- Data types: {df.dtypes.to_dict()}

SAMPLE DATA (first 2 rows):
{json.dumps(sample_data, indent=2, default=str)}

CRITICAL REQUIREMENTS:
1. Use the existing DataFrame variable 'df' (DO NOT create pd.DataFrame or new data)
2. The DataFrame is already loaded and available as 'df'
3. Only use the available columns in the DataFrame
4. Process the data using pandas operations on 'df'
5. Return ONLY the Python code, no explanations
6. DO NOT use import statements - all modules are already available
7. Include ALL variable definitions in your code

AVAILABLE MODULES:
- df: The pandas DataFrame (already loaded)
- pd: pandas module (already imported)
- np: numpy module (already imported)
- Counter: from collections (already imported)
- json: json module (already imported)

CODE REQUIREMENTS:
- Use df.dropna() for missing values
- Use df['column'].str.split(', ') for comma-separated values
- Use .explode() to flatten lists
- Use Counter() for counting (already available)
- Use pd.Series.value_counts() as alternative to Counter
- Limit to top 10-15 items
- Return SIMPLE data structure, not ECharts config
- Define ALL variables before using them in the result

EXAMPLE CODE STRUCTURE:
# Process the data
data_processed = df['column'].dropna()
# ... more processing steps ...
final_data = data_processed.value_counts().head(10)

# Create result with simple data
result = {{
    'success': True,
    'chart_title': 'Clean Title Here',  # NO markdown (#, *, etc.) - just clean text
    'chart_type': 'bar',  # or 'pie', 'line', 'histogram'
    'data': {{
        'labels': final_data.index.tolist(),
        'values': final_data.values.tolist()
    }},
    'data_summary': {{'total_records': len(df), 'processed_columns': [...]}},
    'reasoning': 'Brief explanation'
}}

CHART TITLE REQUIREMENTS:
- Use clean, simple titles (e.g., "Rating Distribution", "City Preferences")
- NO markdown formatting (#, *, _, etc.)
- NO words like "Chart", "Graph", "Visualization"
- 2-4 words maximum
- Title case format

CRITICAL: Return simple data structure, NOT ECharts configuration. The frontend will handle styling.

Generate the complete Python code:"""

        return prompt
    
    def _execute_generated_code(self, code_content: str, df: pd.DataFrame) -> Dict[str, Any]:
        """Execute the generated Python code and return the result."""
        try:
            # Extract code from markdown if present
            if '```python' in code_content:
                code = code_content.split('```python')[1].split('```')[0].strip()
            elif '```' in code_content:
                code = code_content.split('```')[1].strip()
            else:
                code = code_content.strip()
            
            # Clean up the code by removing explanatory text
            lines = code.split('\n')
            code_lines = []
            
            for line in lines:
                # Skip empty lines and pure explanatory text
                if (line.strip() == '' or 
                    line.strip().startswith('Generate') or
                    line.strip().startswith('Output') or
                    line.strip().startswith('ai')):
                    continue
                code_lines.append(line)
            
            code = '\n'.join(code_lines).strip()
            
            # Create a safe execution environment
            safe_globals = {
                'df': df,
                'pd': pd,
                'np': np,
                'json': json,
                'Counter': Counter,
                'len': len,
                'str': str,
                'int': int,
                'float': float,
                'list': list,
                'dict': dict,
                'sorted': sorted,
                'sum': sum,
                'max': max,
                'min': min,
                'round': round,
                'zip': zip,
                'enumerate': enumerate,
                'range': range,
                'print': print,
                'result': None,
                '__builtins__': {
                    'len': len, 'str': str, 'int': int, 'float': float,
                    'list': list, 'dict': dict, 'sorted': sorted,
                    'sum': sum, 'max': max, 'min': min, 'round': round,
                    'zip': zip, 'enumerate': enumerate, 'range': range, 'print': print
                }
            }
            
            # Execute the code
            exec(code, safe_globals)
            
            # Get the result
            if 'result' in safe_globals and safe_globals['result'] is not None:
                result = safe_globals['result']
                if isinstance(result, dict):
                    # Convert numpy types to Python native types for JSON serialization
                    result_str = json.dumps(result, default=str)
                    result = json.loads(result_str)
                    return result
                else:
                     return {
                         'success': False,
                         'error': 'Generated code did not return a valid result dictionary',
                         'chart_title': '',
                         'chart_type': '',
                         'data': {'labels': [], 'values': []},
                         'data_summary': {},
                         'reasoning': ''
                     }
            else:
                return {
                    'success': False,
                    'error': 'No result variable found in generated code',
                    'chart_title': '',
                    'chart_type': '',
                    'data': {'labels': [], 'values': []},
                    'data_summary': {},
                    'reasoning': ''
                }
                
        except Exception as e:
            return {
                'success': False,
                'error': f'Code execution error: {str(e)}',
                'chart_title': '',
                'chart_type': '',
                'data': {'labels': [], 'values': []},
                'data_summary': {},
                'reasoning': ''
            }
    
    def finalize_results_node(self, state: VisualizationProcessingState) -> VisualizationProcessingState:
        """Node 5: Finalize and validate results."""
        try:
            # Clean up chart title - remove markdown formatting
            if state.get("visualization_result", {}).get("chart_title"):
                title = state["visualization_result"]["chart_title"]
                # Remove markdown headers (#, ##, etc.)
                title = re.sub(r'^#+\s*', '', title)
                # Remove markdown formatting (*, _, etc.)
                title = re.sub(r'[*_`]', '', title)
                # Clean up extra whitespace
                title = title.strip()
                # Update the title in the result
                if state.get("visualization_result"):
                    state["visualization_result"]["chart_title"] = title
            
                    
            state["processing_complete"] = True
        except Exception as e:
            state["error"] = f"Finalization error: {str(e)}"
            state["success"] = False
        
        return state
    
    def generate_visualization(
        self, 
        query: str,
        csv_data: str,
        field_info: List[FieldInfo],
        form_title: str
    ) -> Dict[str, Any]:
        """Run the complete workflow and return results."""
        
        # Initialize state
        initial_state = VisualizationProcessingState(
            query=query,
            csv_data=csv_data,
            field_info=field_info,
            form_title=form_title,
            processing_method="structured_llm",
            parsed_dataframe=None,
            data_summary=None,
            column_analysis=None,
            visualization_result=None,
            error=None,
            success=False,
            processing_time=0.0,
            processing_complete=False
        )
        
        try:
            # Run the workflow
            final_state = self.workflow.invoke(initial_state)
            
            return {
                "success": final_state["success"],
                "result": final_state["visualization_result"],
                "error": final_state["error"],
                "processing_time": final_state["processing_time"],
                "usage": final_state.get("usage", {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0})
            }
            
        except Exception as e:
            return {
                "success": False,
                "result": None,
                "error": f"Workflow execution failed: {str(e)}",
                "processing_time": 0.0
            }

# Global workflow instance
workflow_instance = VisualizationWorkflow()