import sys
import pdb
import subprocess
import os
import logging
import tempfile, tarfile
from typing import Optional, List
from io import BytesIO

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
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:

    def __init__(self, 
                 docker_compose_yml:Optional[str]=None,
                 docker_compose_cmd: Optional[List[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"
                    )
        if not docker_compose_cmd:
            self.docker_compose_cmd = ['docker', 'compose']
        else:
            self.docker_compose_cmd = docker_compose_cmd
        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")
        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'Stopping 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 __enter__(self):

        if self.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.code_runner_path,
                    check=True
                    )
            return self

    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
    
    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.code_runner_path,
                    check=True
                    )

class DockerContainerNotFound(RuntimeError):
    pass

class PythonRunnerToolContext():

    def __init__(self, 
                 docker_config: DockerConfig,
                 tmpdir: str, 
                 ignore_dependencies: Optional[List[str]]=None,
                 ignore_unsafe_functions: Optional[List[str]]=None,
                 debug=False
    ):
        """
        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

        # initializations.
        self._tool_node: Optional[ToolNode] = None
        self._tool: Optional[StructuredTool] = None
        self.log = logging.getLogger('PYTHON_RUNNER')

        # 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)
        self.python_runner = AgentRun(
                container_name=container_name,
                cached_dependencies = ['sqlalchemy'],
                install_policy=UVInstallPolicy(),
                cpu_quota=100000,
                default_timeout=100,
                log_level = 'INFO' if self.debug is True else 'WARNING',
                )


    def execute_code(self, code: str):
        return self.python_runner.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_code_to_container(
            self,
            python_code: str,
            dst_path
    ):
        """Copy Python code to the container.
        Args:
            python_code: Python code to copy
            dst_path: Full destination path (incl. file name) inside the container.
        """
        script_path, script_name = os.path.split(dst_path)

        # write the python code to a temporary file.
        temp_script_path = os.path.join(self.tmpdir, script_name)
        with open(temp_script_path, "w") as file:
            file.write(python_code)

        # create a tar stream and add the file to it.
        tar_stream = BytesIO()
        with tarfile.open(fileobj=tar_stream, mode="w") as tar:
            tar.add(temp_script_path, arcname=script_name)
        tar_stream.seek(0)

        # write the file into the container.
        exec_result = self.container.put_archive(path=script_path, data=tar_stream)
        if not exec_result:
            raise RuntimeError(f'Unable to copy code to {dst_path}!')

    def copy_file_to_container(
        self, 
        src_path: str,
        dst_folder: str
    ):
        """Copy Python code to the container.
        Args:
            src_path: Path of the file to copy into container.
            dst_folder: The destination folder inside the container.
        """
        # check if the source file exists.
        if not os.path.isfile(src_path):
            raise RuntimeError(f'{src_path} is not a valid file!')

        # create a tar archive and add the source file.
        tar_stream = BytesIO()
        with tarfile.open(fileobj=tar_stream, mode="w") as tar:
            tar.add(src_path, arcname=os.path.basename(src_path))
        tar_stream.seek(0)

        exec_result = self.container.put_archive(path=dst_folder, data=tar_stream)
        if not exec_result:
            raise RuntimeError(f'Unable to copy {src_path} into container!')

    def copy_file_from_container(
            self, 
            src_path: str,
            dst_folder: Optional[str]=None
            ) -> str:

        # get the archive
        stream, _ = self.container.get_archive(src_path)

        # use temporary folder itself if destination isn't specified.
        if not dst_folder:
            dst_folder = os.path.join(self.tmpdir)

        with tempfile.NamedTemporaryFile(delete=False) as tmpfp:
            try:
                # write the tar stream on to a temporary file ...
                for chunk in stream:
                    tmpfp.write(chunk)
                # close it (the file won't get deleted).
                tmpfp.close()

                # extract it.
                with tarfile.open(tmpfp.name) as tar:
                    tar.extractall(dst_folder)
            finally:
                # delete the tar file eventually.
                os.unlink(tmpfp.name)

        # return the destination file name.
        dst_path = os.path.join(dst_folder, os.path.basename(src_path))
        assert os.path.isfile(dst_path)
        return dst_path

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
