import subprocess
import sys
import os
import tempfile
from fastapi import FastAPI
from pydantic import BaseModel, Field

# Defines the structure of the request body
class ExecutionRequest(BaseModel):
    package_name: str = Field(
        ...,
        example="numpy requests beautifulsoup4",
        description="One or more pip packages to install (space or comma separated)."
    )
    script_content: str = Field(
        ...,
        example="import numpy as np; print(np.array([1, 2, 3]))",
        description="The Python script content to execute."
    )

app = FastAPI()

@app.post("/execute/")
def install_and_execute(request: ExecutionRequest):
    """
    Installs one or more packages into a temporary directory and executes a Python script.
    Supports multiple packages separated by spaces or commas.
    """
    with tempfile.TemporaryDirectory() as temp_dir:
        pip_output = ""
        script_output = ""
        error = ""

        try:
            # 1. Create the install directory
            install_path = os.path.join(temp_dir, "packages")
            os.makedirs(install_path)

            # 2. Split package names by space or comma
            packages = [pkg.strip() for pkg in request.package_name.replace(",", " ").split() if pkg.strip()]

            # 3. Install the packages to the temporary directory
            pip_output = subprocess.check_output(
                [
                    sys.executable,
                    "-m",
                    "pip",
                    "install",
                    *packages,  # multiple packages supported
                    "--target",
                    install_path,
                ],
                stderr=subprocess.STDOUT,
                text=True
            )

            # 4. Create a temporary script file
            script_path = os.path.join(temp_dir, "script.py")
            with open(script_path, 'w') as temp_script:
                temp_script.write(request.script_content)

            # 5. Prepare the environment with custom PYTHONPATH
            exec_env = os.environ.copy()
            exec_env["PYTHONPATH"] = install_path + os.pathsep + exec_env.get("PYTHONPATH", "")

            # 6. Execute the Python script
            script_output = subprocess.check_output(
                [sys.executable, script_path],
                stderr=subprocess.STDOUT,
                text=True,
                env=exec_env
            )

        except subprocess.CalledProcessError as e:
            error = e.output

        return {
            "pip_install_output": pip_output,
            "script_execution_output": script_output,
            "error_output": error
        }