import json
from typing import Optional, Any, List
from .exceptions import (
    LanguageModelBadResponseError,
    LanguageModelTypeError,
    LanguageModelIndexError,
    LanguageModelKeyError
)

from json_repair import repair_json


def _extract_item_from_raw_response(data, sequence):
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except json.JSONDecodeError:
            return data
    current_data = data
    for i, key in enumerate(sequence):
        try:
            if isinstance(current_data, (list, tuple)):
                if not isinstance(key, int):
                    raise LanguageModelTypeError(
                        f"Expected integer index for sequence at position {i}, got {type(key).__name__}"
                    )
                if key < 0 or key >= len(current_data):
                    raise LanguageModelIndexError(
                        f"Index {key} out of range for sequence of length {len(current_data)} at position {i}"
                    )
            elif isinstance(current_data, dict):
                if key not in current_data:
                    raise LanguageModelKeyError(
                        f"Key '{key}' not found in dictionary at position {i}"
                    )
            else:
                raise LanguageModelTypeError(
                    f"Cannot index into {type(current_data).__name__} at position {i}. Full response is: {data} of type {type(data)}. Key sequence is: {sequence}"
                )

            current_data = current_data[key]
        except Exception as e:
            path = " -> ".join(map(str, sequence[: i + 1]))
            
            # Create a safe error message that won't be None
            if "error" in data and data["error"] is not None:
                msg = str(data["error"])
            else:
                msg = f"Error accessing path: {path}. {str(e)}. Full response is: '{data}'"
                
            raise LanguageModelBadResponseError(message=msg, response_json=data)
    if isinstance(current_data, str):
        return current_data.strip()
    else:
        return current_data


class RawResponseHandler:
    """Class to handle raw responses from language models."""

    def __init__(self, key_sequence: list, usage_sequence: Optional[list] = None, reasoning_sequence: Optional[list] = None):
        self.key_sequence = key_sequence
        self.usage_sequence = usage_sequence
        self.reasoning_sequence = reasoning_sequence

    def get_generated_token_string(self, raw_response):
        try:
            return _extract_item_from_raw_response(raw_response, self.key_sequence)
        except (LanguageModelKeyError, LanguageModelIndexError, LanguageModelTypeError, LanguageModelBadResponseError) as e:
            # For non-reasoning models or reasoning models with different response formats,
            # try to extract text directly from common response formats
            if isinstance(raw_response, dict):
                # Responses API format for non-reasoning models
                if 'output' in raw_response and isinstance(raw_response['output'], list):
                    # Try to get first message content
                    if len(raw_response['output']) > 0:
                        item = raw_response['output'][0]
                        if isinstance(item, dict) and 'content' in item:
                            if isinstance(item['content'], list) and len(item['content']) > 0:
                                first_content = item['content'][0]
                                if isinstance(first_content, dict) and 'text' in first_content:
                                    return first_content['text']
                            elif isinstance(item['content'], str):
                                return item['content']
                
                # OpenAI completions format
                if 'choices' in raw_response and isinstance(raw_response['choices'], list) and len(raw_response['choices']) > 0:
                    choice = raw_response['choices'][0]
                    if isinstance(choice, dict):
                        if 'text' in choice:
                            return choice['text']
                        elif 'message' in choice and isinstance(choice['message'], dict) and 'content' in choice['message']:
                            return choice['message']['content']
                
                # Text directly in response
                if 'text' in raw_response:
                    return raw_response['text']
                elif 'content' in raw_response:
                    return raw_response['content']
                
                # Error message - try to return a coherent error for debugging
                if 'message' in raw_response:
                    return f"[ERROR: {raw_response['message']}]"
            
            # If we get a string directly, return it
            if isinstance(raw_response, str):
                return raw_response
            
            # As a last resort, convert the whole response to string
            try:
                return f"[ERROR: Could not extract text. Raw response: {str(raw_response)}]"
            except:
                return "[ERROR: Could not extract text from response]"

    def get_usage_dict(self, raw_response):
        if self.usage_sequence is None:
            return {}
        try:
            return _extract_item_from_raw_response(raw_response, self.usage_sequence)
        except (LanguageModelKeyError, LanguageModelIndexError, LanguageModelTypeError, LanguageModelBadResponseError):
            # For non-reasoning models, try to extract usage from common response formats
            if isinstance(raw_response, dict):
                # Standard OpenAI usage format
                if 'usage' in raw_response:
                    return raw_response['usage']
                
                # Look for nested usage info
                if 'choices' in raw_response and len(raw_response['choices']) > 0:
                    choice = raw_response['choices'][0]
                    if isinstance(choice, dict) and 'usage' in choice:
                        return choice['usage']
            
            # If no usage info found, return empty dict
            return {}

    def get_reasoning_summary(self, raw_response):
        """
        Extract reasoning summary from the model response.
        
        Handles various response structures:
        1. Standard path extraction using self.reasoning_sequence
        2. Direct access to output[0]['summary'] for OpenAI responses
        3. List responses where the first item contains the output structure
        """
        if self.reasoning_sequence is None:
            return None
            
        try:
            # First try the standard extraction path
            summary_data = _extract_item_from_raw_response(raw_response, self.reasoning_sequence)
            
            # If summary_data is a list of dictionaries with 'text' and 'type' fields
            # (as in OpenAI's response format), combine them into a single string
            if isinstance(summary_data, list) and all(isinstance(item, dict) and 'text' in item for item in summary_data):
                return '\n\n'.join(item['text'] for item in summary_data)
            
            return summary_data
        except Exception:
            # Fallback approaches for different response structures
            try:
                # Case 1: Direct dict with 'output' field (common OpenAI format)
                if isinstance(raw_response, dict) and 'output' in raw_response:
                    output = raw_response['output']
                    if isinstance(output, list) and len(output) > 0 and 'summary' in output[0]:
                        summary_data = output[0]['summary']
                        if isinstance(summary_data, list) and all(isinstance(item, dict) and 'text' in item for item in summary_data):
                            return '\n\n'.join(item['text'] for item in summary_data)
                
                # Case 2: List where the first item is a dict with 'output' field
                if isinstance(raw_response, list) and len(raw_response) > 0:
                    first_item = raw_response[0]
                    if isinstance(first_item, dict) and 'output' in first_item:
                        output = first_item['output']
                        if isinstance(output, list) and len(output) > 0 and 'summary' in output[0]:
                            summary_data = output[0]['summary']
                            if isinstance(summary_data, list) and all(isinstance(item, dict) and 'text' in item for item in summary_data):
                                return '\n\n'.join(item['text'] for item in summary_data)
            except Exception:
                pass
            
            return None

    def parse_response(self, raw_response: dict[str, Any]) -> Any:
        """Parses the API response and returns the response text."""

        from edsl.data_transfer_models import EDSLOutput

        generated_token_string = self.get_generated_token_string(raw_response)
        # Ensure generated_token_string is a string before using string methods
        if not isinstance(generated_token_string, str):
            generated_token_string = str(generated_token_string)
        last_newline = generated_token_string.rfind("\n")
        reasoning_summary = self.get_reasoning_summary(raw_response)

        if last_newline == -1:
            # There is no comment
            edsl_dict = {
                "answer": self.convert_answer(generated_token_string),
                "generated_tokens": generated_token_string,
                "comment": None,
                "reasoning_summary": reasoning_summary,
            }
        else:
            edsl_dict = {
                "answer": self.convert_answer(generated_token_string[:last_newline]),
                "comment": generated_token_string[last_newline + 1:].strip(),
                "generated_tokens": generated_token_string,
                "reasoning_summary": reasoning_summary,
            }
        return EDSLOutput(**edsl_dict)

    @staticmethod
    def convert_answer(response_part):
        import json

        response_part = response_part.strip()

        if response_part == "None":
            return None

        repaired = repair_json(response_part)
        if repaired == '""':
            # it was a literal string
            return response_part

        try:
            return json.loads(repaired)
        except json.JSONDecodeError:
            # last resort
            return response_part
