# File: agentq/core/agent/base.py

import logging
import json
import os
from typing import Callable, List, Optional, Tuple, Type, Dict, Any
from dotenv import load_dotenv
from contextlib import contextmanager

from openai import OpenAI
from openai.types.chat import ChatCompletion
import instructor
from instructor import Mode
from langsmith import traceable
from pydantic import BaseModel, ValidationError
from pydantic_settings import BaseSettings

from agentq.utils.function_utils import get_function_schema
from agentq.utils.logger import logger
from agentq.utils.json_cleaner import clean_json

import asyncio

class Settings(BaseSettings):
    OPENROUTER_API_KEY: str
    OPENROUTER_MODEL: str = "openai/gpt-3.5-turbo"

    class Config:
        env_file = ".env"

settings = Settings()

# Загрузка переменных окружения из .env файла
try:
    load_dotenv()
    logger.info("Environment variables loaded successfully.")
except Exception as e:
    logger.error(f"Error loading .env file: {str(e)}")
    raise EnvironmentError("Failed to load environment variables.") from e

@contextmanager
def get_openai_client():
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=settings.OPENROUTER_API_KEY,
    )
    client = instructor.patch(client, mode=Mode.JSON)
    try:
        yield client
    finally:
        # OpenAI client doesn't need explicit closing
        pass

class BaseAgent:
    """
    BaseAgent class for handling interactions with language models and executing tools.
    """

    def __init__(
        self,
        name: str,
        system_prompt: str,
        input_format: Type[BaseModel],
        output_format: Type[BaseModel],
        tools: Optional[List[Tuple[Callable, str]]] = None,
        keep_message_history: bool = True,
        client_type: str = "openrouter",
    ):
        """
        Initialize the BaseAgent.

        :param name: Name of the agent
        :param system_prompt: System prompt for the language model
        :param input_format: Pydantic model for input validation
        :param output_format: Pydantic model for output validation
        :param tools: List of tools (functions) available to the agent
        :param keep_message_history: Whether to keep message history
        :param client_type: Type of client to use (e.g., "openrouter")
        """
        self.agent_name = name
        self.system_prompt = system_prompt
        self.input_format = input_format
        self.output_format = output_format
        self.keep_message_history = keep_message_history

        self.default_model = settings.OPENROUTER_MODEL

        self._initialize_tools(tools)

        if self.system_prompt:
            self._initialize_messages()
            logger.debug("System prompt initialized in message history.")

    def _initialize_tools(self, tools: Optional[List[Tuple[Callable, str]]]):
        """Initialize tools for the agent."""
        self.tools_list: List[Any] = []
        self.executable_functions_list: Dict[str, Callable] = {}
        if tools:
            for func, func_desc in tools:
                if func.__name__ in self.executable_functions_list:
                    logger.warning(f"Duplicate function name detected: {func.__name__}. Overwriting previous definition.")
                function_schema = get_function_schema(func, description=func_desc)
                self.tools_list.append(function_schema)
                self.executable_functions_list[func.__name__] = func
                logger.debug(f"Registered tool: {func.__name__} - {func_desc}")
        else:
            logger.info("No tools provided for initialization.")

    def _initialize_messages(self):
        """Initialize the message history with the system prompt."""
        self.messages: List[Dict[str, Any]] = [{"role": "system", "content": self.system_prompt}]
        logger.debug("Message history initialized with system prompt.")

    @traceable(run_type="chain", name="agent_run")
    async def run(
        self,
        input_data: BaseModel,
        screenshot: Optional[str] = None,
        session_id: Optional[str] = None,
        model: Optional[str] = None,
    ) -> BaseModel:
        """
        Run the agent with the given input data.

        :param input_data: Input data conforming to the input_format
        :param screenshot: Optional screenshot URL
        :param session_id: Optional session ID for tracking
        :param model: Optional model override
        :return: Output data conforming to the output_format
        """
        if not isinstance(input_data, self.input_format):
            logger.error(f"Invalid input data type: expected {self.input_format.__name__}, got {type(input_data).__name__}")
            raise ValueError(f"Input data must be of type {self.input_format.__name__}")

        if not self.keep_message_history:
            self._initialize_messages()
            logger.debug("Message history re-initialized due to keep_message_history=False.")

        self._append_input_message(input_data, screenshot)

        model_to_use = model or self.default_model
        logger.info(f"Using model: {model_to_use}")

        try:
            with get_openai_client() as client:
                response = client.chat.completions.create(
                    model=model_to_use,
                    messages=self.messages,
                    max_retries=4
                )
            logger.debug(f"Received response from model: {response}")
        except Exception as e:
            logger.error(f"Error calling OpenAI API: {str(e)}")
            raise

        return await self._process_response(response, session_id)

    def _append_input_message(self, input_data: BaseModel, screenshot: Optional[str]):
        """Append input data to the message history."""
        try:
            content = input_data.model_dump_json(exclude={"current_page_dom", "current_page_url"})
            logger.debug(f"Input content (JSON): {content}")

            if screenshot:
                self.messages.append({
                    "role": "user",
                    "content": [
                        {"type": "text", "text": content},
                        {"type": "image_url", "image_url": {"url": screenshot}},
                    ],
                })
                logger.debug("Appended user message with screenshot.")
            else:
                self.messages.append({"role": "user", "content": content})
                logger.debug("Appended user message without screenshot.")

            if hasattr(input_data, "current_page_dom") and hasattr(input_data, "current_page_url"):
                current_page_content = (
                    f"Current page URL:\n{input_data.current_page_url}\n\n"
                    f"Current page DOM:\n{input_data.current_page_dom}"
                )
                self.messages.append({
                    "role": "user",
                    "content": current_page_content,
                })
                logger.debug("Appended current page URL and DOM to user messages.")
        except Exception as e:
            logger.error(f"Error appending input message: {str(e)}")
            raise

    async def _process_response(self, response: ChatCompletion, session_id: Optional[str]) -> BaseModel:
        """
        Process the response from the language model.

        :param response: The response from the OpenAI API
        :param session_id: Optional session ID for tracking
        :return: Processed and validated response as a BaseModel
        """
        try:
            response_content = response.choices[0].message.content
            logger.debug(f"Raw response content: {response_content}")

            cleaned_content = clean_json(response_content)
            logger.debug(f"Cleaned JSON content: {cleaned_content}")

            parsed_json = json.loads(cleaned_content)
            logger.debug(f"Parsed JSON: {parsed_json}")

            self._process_parsed_json(parsed_json)

            logger.debug(f"Processed JSON before validation: {parsed_json}")

            parsed_response = self.output_format.parse_obj(parsed_json)
            assert isinstance(parsed_response, self.output_format)
            logger.info("Successfully parsed and validated response.")
            return parsed_response
        except json.JSONDecodeError as e:
            logger.error(f"JSON decode error: {str(e)}")
            logger.error(f"Problematic JSON: {response_content}")
            raise ValueError(f"Response is not a valid JSON: {str(e)}") from e
        except ValidationError as e:
            logger.error(f"Validation error: {str(e)}")
            raise TypeError(f"Failed to validate response as {self.output_format.__name__}: {str(e)}") from e
        except Exception as e:
            logger.error(f"Unexpected error processing response: {str(e)}")
            logger.error(f"Raw response content: {response_content}")
            raise

    def _process_parsed_json(self, parsed_json: Dict[str, Any]):
        """
        Process and structure the parsed JSON data.

        This method performs the following operations:
        1. Processes the 'plan' field, ensuring each task has 'url' and 'result' fields.
        2. Converts 'next_task' to a dictionary if it's a string.
        3. Sets default values for various fields.
        4. Processes the 'next_task_actions' field.

        :param parsed_json: The parsed JSON data to process
        """
        # Обработка плана задач
        for task in parsed_json.get('plan', []):
            task.setdefault('url', None)
            task.setdefault('result', None)
            logger.debug(f"Processed task: {task}")

        # Обработка следующей задачи
        if isinstance(parsed_json.get('next_task'), str):
            parsed_json['next_task'] = {'description': parsed_json['next_task']}
            logger.debug(f"Converted next_task from string to dict: {parsed_json['next_task']}")

        # Установка значений по умолчанию
        parsed_json.setdefault('thought', '')
        parsed_json.setdefault('plan', [])
        parsed_json.setdefault('next_task', {})
        parsed_json.setdefault('next_task_actions', [])
        parsed_json.setdefault('is_complete', False)
        parsed_json.setdefault('final_response', None)

        self._process_next_task_actions(parsed_json)

    def _process_next_task_actions(self, parsed_json: Dict[str, Any]):
        """
        Process and structure the next_task_actions in the parsed JSON.

        This method handles different formats of actions, converting them into a standardized structure.
        It specifically handles 'GOTO_URL' actions and provides a framework for adding more action types.

        :param parsed_json: The parsed JSON data containing next_task_actions to process
        """
        for i, action in enumerate(parsed_json.get('next_task_actions', [])):
            if isinstance(action, str):
                try:
                    action_type, details = action.split('[', 1)
                    details = details.rstrip(']')
                    if action_type == 'GOTO_URL':
                        url, timeout = details.split(',', 1)
                        parsed_json['next_task_actions'][i] = {
                            'type': 'GOTO_URL',
                            'website': url.strip(),
                            'timeout': int(timeout.strip())
                        }
                        logger.debug(f"Processed action GOTO_URL: {parsed_json['next_task_actions'][i]}")
                    # Добавьте другие типы действий по мере необходимости
                    else:
                        parsed_json['next_task_actions'][i] = {
                            'type': action_type,
                            'details': details.strip()
                        }
                        logger.debug(f"Processed action {action_type}: {details.strip()}")
                except ValueError:
                    logger.warning(f"Invalid action format: {action}")
                    parsed_json['next_task_actions'][i] = {'type': 'UNKNOWN', 'original': action}
            else:
                logger.warning(f"Unexpected action type: {type(action).__name__} in action {action}")
                parsed_json['next_task_actions'][i] = {'type': 'UNKNOWN', 'original': action}

    async def _append_tool_response(self, tool_call: Any):
        """
        Append the response from a tool call to the message history.

        :param tool_call: The tool call object
        """
        function_name = tool_call.function.name
        function_to_call = self.executable_functions_list.get(function_name)
        if not function_to_call:
            logger.error(f"Attempted to call unknown function: {function_name}")
            raise ValueError(f"Unknown function: {function_name}")

        try:
            function_args = json.loads(tool_call.function.arguments)
            logger.debug(f"Function arguments for {function_name}: {function_args}")

            function_response = await function_to_call(**function_args)
            self.messages.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": function_name,
                "content": str(function_response),
            })
            logger.info(f"Appended tool response for function: {function_name}")
        except json.JSONDecodeError as e:
            logger.error(f"Error parsing function arguments for {function_name}: {str(e)}")
            raise ValueError(f"Invalid function arguments for {function_name}") from e
        except Exception as e:
            logger.error(f"Error occurred calling the tool {function_name}: {str(e)}")
            self.messages.append({
                "tool_call_id": tool_call.id,
                "role": "tool",
                "name": function_name,
                "content": "The tool responded with an error. Please try again with a different tool or modify the parameters of the tool.",
            })
            logger.debug(f"Appended error message for tool {function_name}.")