import sys
import pdb
import subprocess
import os
from typing import Optional, List
import abc

import docker
import docker.errors
from docker import DockerClient
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import StructuredTool
from langchain.output_parsers import PydanticOutputParser
from langchain_core.tools import BaseTool, ArgsSchema, BaseToolkit
import yaml

from agentrun_plus import AgentRun, UVInstallPolicy, InstallPolicy
from loguru import logger

class PythonToolInput(BaseModel):
    python_code: str = Field(description=" Python code to run provided as a string.")
    artifacts_abs_paths: List[str] = Field(description="""
    If the python code is generating a artifact (such as a plot), specify the
    absolute path of the files containing the artifacts.
    """)

class PythonToolOutput(BaseModel):
    stdout: str = Field(description="Text output of the Python Program")
    user_artifacts_abs_paths: List[str] = Field(description="""
    Local paths of artifacts (if any) generated by the Python program, corresponding to the input `artifacts_abs_path` parameter. This is only meant for user display. You can continue to use the original artifact locations for further use in downstream tool calls.
    """)

PYTHON_TOOL_DESCRIPTION="""
Runs the specified python code provided and returns the standard output
text. This tool doesn't have any display capabilities, and hence, if you
wish to generate a plot, store the plot/figure into a PNG file and specify
its absolute path in the `artifacts_abs_paths` parameter. 

Your code will run in a docker environment with a basic python base image
(`python:3.12.2-slim-bullseye`) and the following packages installed:

{packages}

So, restrict your code accordingly.

Always store artifacts such as plots in the `/home/pythonuser/artifacts`
folder. The `artifacts_abs_paths` parameter allows you to generate and pass
multiple artifact paths. Pass an empty list if no artifacts are generated by
the program. 

The tool returns modified locations that you can pass to the user, but you
can continue to refer to the original locations (in `/home/pythonuser/artifacts`) for
referencing in downstream invocations of the tool.

The tool output follows this schema:
{output_schema}
"""
class DockerConfig:
    """
    Encapsulates a Docker configuration specified by a docker_compose.yml file.
    """

    def __init__(self, 
                 docker_compose_yml:Optional[str]=None
    ):

        if not docker_compose_yml:
            docker_compose_yml = os.path.join(
                    os.path.dirname(os.path.abspath(__file__)),
                    "code_runner", 
                    "docker-compose.yml"
                    )
        with open(docker_compose_yml, 'rt') as f:
            self.docker_yml = yaml.load(f, Loader=yaml.Loader)
        self.code_runner_path = os.path.dirname(docker_compose_yml)

        self.log = logger.bind(name="Docker-Config")
        self.log.remove()
        self.log.add(
                sys.stderr,
                format="{time} {level} {message}",
                level="INFO",
                enqueue=True
        )
        self.running_on_enter = False

    @staticmethod
    def is_docker_installed():
        try:
            client = docker.from_env()
            _ = client.version()
            return True
        except docker.errors.DockerException:
            return False

    def is_container_running(self):

        client = docker.from_env()

        for _, config in self.docker_yml['services'].items():
            container_name = config['container_name']
            try:
                container = client.containers.get(container_name)
                match container.status:
                    case "running":
                        # if the container is running we are good ...
                        pass
                    case "exited":
                        # if the container was stopped remove it before continuing.
                        self.log.info(f'Removing container {container_name} ...')
                        container.remove()
                    case _:
                        # otherwise, assume it can started again.
                        return False
            except docker.errors.NotFound:
                self.log.error(f"Container '{container_name}' not found.")
                return False
            except docker.errors.DockerException as e:
                self.log.error(f"Error: {e}")
                return False

        return True

    def get_container_name(self, service_name: str):
        return self.docker_yml['services'][service_name]['container_name']

    def get_requirements_txt(self):
        reqs_file = os.path.join(self.code_runner_path, 'requirements.txt')
        assert os.path.isfile(reqs_file)
        return reqs_file
    

class DockerRunner:
    """
    A context manager for starting and stopping the docker container defined by
    DockerConfig.
    """
    def __init__(
            self, 
            docker_config: DockerConfig,
            docker_compose_cmd: Optional[List[str]]=None
        ):
        self.docker_config = docker_config
        if docker_compose_cmd is None:
            self.docker_compose_cmd = [ 'docker', 'compose' ]
        else:
            self.docker_compose_cmd = docker_compose_cmd
        self.log = logger.bind(name="Docker-Runner")
        self.log.remove()
        self.log.add(
                sys.stderr,
                format="{time} {level} {message}",
                level="INFO",
                enqueue=True
        )
        self.running_on_enter = False

    def __enter__(self):

        if self.docker_config.is_container_running():
            self.log.info(f'Docker container is already running ...')
            self.running_on_enter = True
            return self
        else:
            self.running_on_enter = False
            self.log.info('Starting Docker Container ...')
            subprocess.run(
                    self.docker_compose_cmd + ['up', '--build', '-d'], 
                    cwd=self.docker_config.code_runner_path,
                    check=True
                    )
            return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # don't stop the docker container if it was running on entry.
        if not self.running_on_enter:
            self.log.info('Stopping Python Container ...')
            subprocess.run(
                    self.docker_compose_cmd + ['down'],
                    cwd=self.docker_config.code_runner_path,
                    check=True
                    )

class DockerContainerNotFound(RuntimeError):
    pass

class PythonAppPolicy(abc.ABC):
    """
    This policy class allows an application to customize the "PythonRunner"
    tool for their use case. The code that an LLM produces will run inside a
    Docker container with a few specific requirements that are described in the
    docstring for PythonRunnerToolContext.

    You can work with the following assumptions.

    - `/home/pythonuser`: This is the working folder for any code executed by
      the LLM.
    - `/home/pythonuser/src`: Hosts any source files needed to expose the any
      objects to the LLM. 
        - As an example, a SQL agent will host a file called "database.py"
          inside the "src/" folder from which an LLM can import the engine
          object using a line such as `from src.database import engine`. Any
          additional files needed by the application (such as an sqlite
          database), can be copied to this folder by overriding the
          `files_to_copy` method.
    - `/home/pythonuser/artifacts`: This folder inside the container contains
      artifacts generated by the LLM code (such as PNG plots).
    """
    def files_to_copy(self) -> List[str]:
        """
        Specify a list of files to copy into the working directory of the
        Docker Container.
        """
        return []

class PythonRunnerToolContext():
    """
    --------------------
    Docker Requirements:
    --------------------
    - The docker container must be setup with a user called "pythonuser" with a
      home folder called "/home/pythonuser" which is also the current working
      directory.
    - The docker container must have a service called "python runner" that must
      be setup prior to creating this context.
    """

    def __init__(self, 
                 app_policy: PythonAppPolicy,
                 docker_config: DockerConfig,
                 tmpdir: str,
                 ignore_dependencies: Optional[List[str]]=None,
                 ignore_unsafe_functions: Optional[List[str]]=None,
                 debug=False,
                 install_policy: Optional[InstallPolicy]=None
    ):
        """
        Initialize the PythonRunnerTool.

        Args:
            docker_config: A DockerConfig instance that starts the docker
                image using a context manager.
            tempdir: A temporary folder to store temporary files (to be deleted
                on program exit). Use the tempfile.TemporaryDirectory context
                manager for this.
        """
        # latch parameters.
        self.debug = debug
        self.client: DockerClient = docker.from_env()
        self.ignore_dependencies = ignore_dependencies
        self.ignore_unsafe_functions = ignore_unsafe_functions
        self.tmpdir = tmpdir
        self.app_policy = app_policy

        # initializations.
        self._tool_node: Optional[ToolNode] = None
        self._tool: Optional[StructuredTool] = None
        self.log = logger.bind(name="Python-Runner")
        self.log.remove()
        self.log.add(
                sys.stderr,
                format="{time} {level} {message}",
                level="INFO",
                enqueue=True
        )

        # check if Docker is running by the time this context is created, and
        # intialize AgentRun.
        if not docker_config.is_container_running():
            raise DockerContainerNotFound(f"Docker container service 'python_runner' isn't running!")
        self.docker_config = docker_config
        container_name = self.docker_config.get_container_name("python_runner")
        self.container = self.client.containers.get(container_name)
        if install_policy is None:
            install_policy = UVInstallPolicy()
        self.agent_run = AgentRun(
                container_name=container_name,
                cached_dependencies = ['sqlalchemy'],
                install_policy=install_policy,
                cpu_quota=100000,
                default_timeout=100,
                log_level = 'INFO' if self.debug is True else 'WARNING',
                )

        # copy any files the application needs.
        for file in self.app_policy.files_to_copy():
            self.agent_run.copy_file_to_container(file, '/home/pythonuser/src')

    def copy_code_to_container(
            self,
            python_code: str,
            dest_file_path: str):
        self.agent_run.copy_code_to_container(python_code, dest_file_path)

    def execute_code(self, code: str):
        return self.agent_run.execute_code_in_container(
                code, 
                ignore_dependencies=self.ignore_dependencies,
                ignore_unsafe_functions=self.ignore_unsafe_functions
        )

    def configure(self, config: RunnableConfig):
        config['configurable']['python_runner'] = self  # pyright: ignore[reportTypedDictNotRequiredAccess]

    def get_requirements_txt(
            self,
            fmt="markdown"
    ) -> str:
        requirements_txt = self.docker_config.get_requirements_txt()
        lines: List[str] = []
        with open(requirements_txt, 'rt') as f:
            lines = f.readlines()
        match fmt:
            case 'markdown':
                lines = [ f'  - {line.rstrip()}' for line in lines ]
                return '\n'.join(lines)
            case _: 
                raise RuntimeError(f'Unknown format: {format}!')

    def copy_file_from_container(
            self, 
            src_path: str,
            dst_folder: Optional[str]=None
            ) -> str:
        if dst_folder is None:
            dst_folder = self.tmpdir
        return self.agent_run.copy_file_from_container(src_path, dst_folder)

class PythonRunnerTool(BaseTool):

    # overrides from BaseTool
    name: str = "run_python_tool"
    args_schema: ArgsSchema | None = PythonToolInput 
    description: str = ""   # fill this in during construction.
    return_direct: bool = True

    # our configuration.
    config: PythonRunnerToolContext
    
    def _run(
            self,
            python_code: str,
            artifacts_abs_paths: List[str]
    ):
        python_runner = self.config
        output = python_runner.execute_code(python_code)
        local_files = []
        for cont_path in artifacts_abs_paths:
            try:
                # copy the file locally
                local_files.append(python_runner.copy_file_from_container(cont_path))
            except docker.errors.NotFound: 
                raise RuntimeError(f"{cont_path} is inaccessible. The program may have failed - here's the output:\n {output}")
        return PythonToolOutput(
                stdout=output,
                user_artifacts_abs_paths=local_files
        )

class PythonRunnerToolkit(BaseToolkit):
    tools: List[BaseTool]

    @classmethod
    def from_context(
            cls,
            config: PythonRunnerToolContext = Field(exclude=True)
    ):
        output_schema = PydanticOutputParser(pydantic_object=PythonToolOutput).get_output_jsonschema()
        packages = config.get_requirements_txt(fmt='markdown')
        return cls(
                tools=[
                    PythonRunnerTool(
                        description=PYTHON_TOOL_DESCRIPTION.format(
                            output_schema=output_schema,
                            packages=packages,
                        ),
                        config=config
                    )
                ]
            )

    def get_tools(self) -> List[BaseTool]:
        return self.tools
