"""
ResultParser - parses AIGC responses and merges results with original data.
"""

import json
import logging
from typing import Union, Dict, Any
import pandas as pd

from .exceptions import ResultParseError

logger = logging.getLogger(__name__)


class ResultParser:
    """
    Result parser that handles AIGC responses and merges them with original data.
    """
    
    @staticmethod
    def parse_prompt_response(response: dict, include_reasoning: bool = False) -> dict:
        """
        Parse Prompt application response.
        
        Args:
            response: API response dictionary
            include_reasoning: Whether to include reasoning chain
            
        Returns:
            dict: Parsed result containing:
                - success: bool
                - result: str or dict
                - reasoning: str (if include_reasoning=True)
                - judge_time: str (request timestamp)
                - log_id: str (request log ID)
                - error: str (if failed)
        """
        try:
            # Extract metadata
            judge_time = response.get('_judge_time', '')
            log_id = response.get('_log_id', '')
            
            # Check if response is successful
            if response.get('code') != '000':
                error_msg = response.get('message', 'Unknown error')
                logger.error(f"Prompt application failed: {error_msg}")
                return {
                    'success': False,
                    'result': None,
                    'reasoning': None,
                    'judge_time': judge_time,
                    'log_id': log_id,
                    'error': error_msg
                }
            
            # Extract text from response
            data = response.get('data', {})
            text = data.get('text', '')
            
            # Extract reasoning if requested
            reasoning = None
            if include_reasoning:
                reasoning = data.get('reasoning_content', '')
            
            # Try to parse as JSON, otherwise treat as string
            if ResultParser.is_json(text):
                try:
                    result = json.loads(text)
                except json.JSONDecodeError:
                    logger.warning(f"Failed to parse JSON, treating as string: {text}")
                    result = text
            else:
                result = text
            
            return {
                'success': True,
                'result': result,
                'reasoning': reasoning,
                'judge_time': judge_time,
                'log_id': log_id,
                'error': None
            }
            
        except Exception as e:
            logger.error(f"Error parsing prompt response: {str(e)}")
            return {
                'success': False,
                'result': None,
                'reasoning': None,
                'judge_time': response.get('_judge_time', ''),
                'log_id': response.get('_log_id', ''),
                'error': str(e)
            }
    
    @staticmethod
    def parse_agent_response(response: dict, include_reasoning: bool = False) -> dict:
        """
        Parse Agent application response.
        
        Args:
            response: API response dictionary
            include_reasoning: Whether to include reasoning chain
            
        Returns:
            dict: Parsed result containing:
                - success: bool
                - result: str or dict
                - reasoning: str (if include_reasoning=True)
                - judge_time: str (request timestamp)
                - log_id: str (conversation ID)
                - error: str (if failed)
        """
        try:
            # Extract metadata
            judge_time = response.get('_judge_time', '')
            log_id = response.get('_log_id', '')
            
            # Check if response is successful
            if response.get('code') != 'success':
                error_msg = response.get('message', 'Unknown error')
                logger.error(f"Agent application failed: {error_msg}")
                return {
                    'success': False,
                    'result': None,
                    'reasoning': None,
                    'judge_time': judge_time,
                    'log_id': log_id,
                    'error': error_msg
                }
            
            # Extract content from response
            data = response.get('data', [])
            if not data or len(data) == 0:
                return {
                    'success': False,
                    'result': None,
                    'reasoning': None,
                    'judge_time': judge_time,
                    'log_id': log_id,
                    'error': 'Empty response data'
                }
            
            content = data[-1].get('content', '')
            
            # Extract reasoning if requested
            reasoning = None
            if include_reasoning:
                reasoning = data[0].get('reasoning_content', '')
            
            # Try to parse as JSON, otherwise treat as string
            if ResultParser.is_json(content):
                try:
                    result = json.loads(content)
                except json.JSONDecodeError:
                    logger.warning(f"Failed to parse JSON, treating as string: {content}")
                    result = content
            else:
                result = content
            
            return {
                'success': True,
                'result': result,
                'reasoning': reasoning,
                'judge_time': judge_time,
                'log_id': log_id,
                'error': None
            }
            
        except Exception as e:
            logger.error(f"Error parsing agent response: {str(e)}")
            return {
                'success': False,
                'result': None,
                'reasoning': None,
                'judge_time': response.get('_judge_time', ''),
                'log_id': response.get('_log_id', ''),
                'error': str(e)
            }

    
    @staticmethod
    def is_json(text: str) -> bool:
        """
        Check if a string is valid JSON.
        
        Args:
            text: String to check
            
        Returns:
            bool: True if valid JSON, False otherwise
        """
        if not text or not isinstance(text, str):
            return False
        
        text = text.strip()
        if not text:
            return False
        
        # Check if it starts with { or [
        if not (text.startswith('{') or text.startswith('[')):
            return False
        
        try:
            json.loads(text)
            return True
        except (json.JSONDecodeError, ValueError):
            return False

    
    @staticmethod
    def merge_result_to_row(row: pd.Series, result: dict) -> pd.Series:
        """
        Merge parsed result to original data row.
        
        Args:
            row: Original data row (pandas Series)
            result: Parsed result dictionary from parse_prompt_response or parse_agent_response
            
        Returns:
            pd.Series: Merged data row with new columns added at the end
        """
        # Create a copy of the row to avoid modifying the original
        merged_row = row.copy()
        
        # Always add judge_time and log_id metadata
        merged_row['judge_time'] = result.get('judge_time', '')
        merged_row['log_id'] = result.get('log_id', '')
        
        # Always add _error column to maintain consistent schema
        # If parsing failed, set error message; otherwise set None
        if not result['success']:
            merged_row['_error'] = result['error']
            return merged_row
        else:
            # Success case: add _error column with None value
            merged_row['_error'] = None
        
        # Get the result data
        result_data = result['result']
        
        # Handle JSON result (dict)
        if isinstance(result_data, dict):
            # Add each key-value pair as a new column
            for key, value in result_data.items():
                # Handle column name conflicts by adding suffix
                column_name = key
                suffix = 1
                while column_name in merged_row.index:
                    column_name = f"{key}_{suffix}"
                    suffix += 1
                
                merged_row[column_name] = value
        
        # Handle string result
        else:
            # Create "result" column for string results
            column_name = 'result'
            suffix = 1
            while column_name in merged_row.index:
                column_name = f"result_{suffix}"
                suffix += 1
            
            merged_row[column_name] = str(result_data)
        
        # Add reasoning content if present
        if result.get('reasoning') is not None:
            column_name = 'reasoning_content'
            suffix = 1
            while column_name in merged_row.index:
                column_name = f"reasoning_content_{suffix}"
                suffix += 1
            
            merged_row[column_name] = result['reasoning']
        
        return merged_row
