import io
from contextlib import redirect_stderr, redirect_stdout
from typing import Any, Dict, Type

import numpy as np
import pandas as pd
from crewai.tools import BaseTool
from crewai_tools import CodeInterpreterTool
from crewai_tools.tools.code_interpreter_tool.code_interpreter_tool import SandboxPython
from pydantic import BaseModel, Field


class ExtractSchemaInput(BaseModel):
    """Input schema for ExtractSchema."""
    file_path: str = Field(
        ..., description="Path to the local file from which to extract the schema, " \
        "relative to the agent's working directory."
    )
    # This file path is relative to ..? (find working dir of agent)

class ExtractSchema(BaseTool):
    name: str = "ExtractSchema"
    description: str = (
        "This tool is useful for extracting schema of tabular data, which is important to understand"
        "the structure of a dataset for analysis."
    )
    args_schema: Type[BaseModel] = ExtractSchemaInput

    def _run(self, file_path: str) -> Dict[str, Any]:
        # Implementation goes here
        df = pd.read_csv(file_path)
        return {
            "shape": df.shape,
            "dtypes": df.dtypes.to_dict(),
            "columns": df.columns.tolist(),
            "numeric_cols": df.select_dtypes(include=[np.number]).columns.tolist(),
            "categorical_cols": df.select_dtypes(include=['object']).columns.tolist()
        }


class ResultAwareCodeInterpreterTool(CodeInterpreterTool):
    """Returns `result` when present and falls back to captured stdout/stderr."""

    @staticmethod
    def _resolve_execution_output(
        exec_locals: Dict[str, Any],
        stdout_buffer: io.StringIO,
        stderr_buffer: io.StringIO,
    ) -> str:
        result = exec_locals.get("result")
        if result is not None:
            return str(result)

        stdout_text = stdout_buffer.getvalue().strip()
        stderr_text = stderr_buffer.getvalue().strip()

        if stdout_text and stderr_text:
            return f"{stdout_text}\n\n[stderr]\n{stderr_text}"
        if stdout_text:
            return stdout_text
        if stderr_text:
            return stderr_text
        return "Code executed successfully, but it did not set `result` or print output."

    @staticmethod
    def run_code_in_restricted_sandbox(code: str) -> str:
        stdout_buffer = io.StringIO()
        stderr_buffer = io.StringIO()
        exec_locals: Dict[str, Any] = {}

        try:
            with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
                SandboxPython.exec(code=code, locals_=exec_locals)
        except Exception as e:
            return f"An error occurred: {e!s}"

        return ResultAwareCodeInterpreterTool._resolve_execution_output(
            exec_locals,
            stdout_buffer,
            stderr_buffer,
        )

    @staticmethod
    def run_code_unsafe(code: str, libraries_used: list[str]) -> str:
        stdout_buffer = io.StringIO()
        stderr_buffer = io.StringIO()

        for library in libraries_used:
            import os

            os.system(f"pip install {library}")  # noqa: S605

        try:
            exec_locals: Dict[str, Any] = {}
            with redirect_stdout(stdout_buffer), redirect_stderr(stderr_buffer):
                exec(code, {}, exec_locals)  # noqa: S102
        except Exception as e:
            return f"An error occurred: {e!s}"

        return ResultAwareCodeInterpreterTool._resolve_execution_output(
            exec_locals,
            stdout_buffer,
            stderr_buffer,
        )