import aiohttp
import io
import os
import shlex
import shutil
import subprocess
import toml
import zipfile

import sys

from loguru import logger
from pathlib import Path

from vr_cli.utils.config import VOICERUN_DIR

# Global directories for module builds and virtual environments
BUILDS_DIR = VOICERUN_DIR / "builds"
VENV_DIR = VOICERUN_DIR / "venv"

# Module-level verbose flag
_verbose = False


def configure_logging(verbose: bool = False) -> None:
    """Configure logging level. Set verbose=True for debug output."""
    global _verbose
    _verbose = verbose
    logger.remove()
    level = "DEBUG" if verbose else "WARNING"
    logger.add(sys.stderr, level=level)


def _ensure_venv() -> Path:
    """
    Ensure a virtual environment exists in the global .voicerun directory.
    Creates one at ~/.voicerun/venv if it doesn't exist.
    Returns the path to the venv.
    """
    venv_python = VENV_DIR / "bin" / "python"

    if not VENV_DIR.exists():
        logger.debug(f"Creating virtual environment at {VENV_DIR}")
        subprocess.run(
            ["uv", "venv", str(VENV_DIR)],
            check=True,
            capture_output=not _verbose,
        )
        # Install test dependencies in the new venv
        logger.debug("Installing test dependencies in venv")
        subprocess.run(
            ["uv", "pip", "install", "--python", str(venv_python), "pytest", "pytest-asyncio"],
            check=True,
            capture_output=not _verbose,
        )
    else:
        logger.debug(f"Using existing virtual environment at {VENV_DIR}")

    return VENV_DIR


async def _make_request(endpoint: str, method: str = "GET", data: dict = None) -> dict:
    """Make a request to the VoiceRun API."""
    async with aiohttp.ClientSession() as session:
        api_base_url = os.environ.get("VOICERUN_API_URL", "https://api.voicerun.com")
        api_key = os.environ.get("VOICERUN_API_KEY", "")

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }
                    
        async with session.request(method, f"{api_base_url}{endpoint}", headers=headers, data=data) as response:
            response.raise_for_status()

            if response.status == 200 or response.status == 201:
                response_json = await response.json()

                return response_json["data"]
            else:
                raise ValueError(f"Request to {endpoint} failed with status {response.status}")

async def _get_secrets_for_agent(agent_id: str, environment: str):
    """
    Get the variables for an agent.
    """
    try:
        response = await _make_request(f"/v1/agents/{agent_id}/run?environment={environment}")
        variables = response.get("secrets", {})

        # Write secrets to a file that can be sourced by shell scripts.
        # This allows child processes (like pytest) to access these secrets.
        variables_file = Path("./.variables.env")
        with variables_file.open("w") as f:
            for key, value in variables.items():
                # Use shlex.quote to properly escape the value for shell safety
                escaped_value = shlex.quote(str(value))
                f.write(f"export {key}={escaped_value}\n")
        
        logger.debug(f"Wrote {len(variables)} variables to {variables_file}")
    except Exception as e:
        logger.error(f"Error getting variables for agent {agent_id}: {e}")

async def _load_module_from_local_path(local_path: str, module_dir: Path) -> tuple[bytes, bool]:
    """
    Load module code from a local file path.
    Supports:
    - Single Python file (will be saved as handler.py)
    - Directory containing Python files (will be copied)
    - ZIP file (will be extracted)
    """
    local_path_obj = Path(local_path).resolve()
    
    if not local_path_obj.exists():
        raise ValueError(f"Local module path does not exist: {local_path}")
    
    # If it's a directory, copy all Python files and other relevant files
    if local_path_obj.is_dir():
        logger.debug(f"Loading module from local directory: {local_path}")
        
        # Copy all Python files
        for py_file in local_path_obj.glob("*.py"):
            shutil.copy2(py_file, module_dir / py_file.name)
        
        # Copy pyproject.toml and requirements.txt if they exist
        for config_file in ["pyproject.toml", "requirements.txt"]:
            config_path = local_path_obj / config_file
            if config_path.exists():
                shutil.copy2(config_path, module_dir / config_file)
        
        # Return empty bytes and False (not a zip) since we've already copied files
        return b"", False
    
    # If it's a file, check if it's a ZIP or a Python file
    if local_path_obj.is_file():
        content = local_path_obj.read_bytes()
        
        # Detect if content is a ZIP archive (by magic bytes)
        is_zip = (len(content) >= 2 and content[:2] == b'PK')
        
        if is_zip:
            logger.debug(f"Loading module from local ZIP file: {local_path}")
            with zipfile.ZipFile(io.BytesIO(content)) as zip_ref:
                zip_ref.extractall(module_dir)
        else:
            # Treat as a single Python file
            logger.debug(f"Loading module from local Python file: {local_path}")
            try:
                code_str = content.decode("utf-8")
            except UnicodeDecodeError:
                raise ValueError(f"Local file is not valid UTF-8 text and is not a zip file: {local_path}")
            (module_dir / "handler.py").write_text(code_str, encoding="utf-8")
        
        return content, is_zip
    
    raise ValueError(f"Local module path is neither a file nor a directory: {local_path}")


def _get_dependencies_from_pyproject(pyproject_path: Path) -> list[str]:
    dependencies: list[str] = []

    try:
        pyproject_dict = toml.loads(pyproject_path.read_text(encoding="utf-8"))
        pyproject_dependencies = pyproject_dict.get("project", {}).get("dependencies", [])
        if isinstance(pyproject_dependencies, list):
            dependencies = pyproject_dependencies
    except Exception as e:
        logger.warning(f"Could not read dependencies from user pyproject.toml: {e}")
      
    return dependencies

def _get_dependencies_from_requirements(req_path: Path) -> list[str]:
    dependencies: list[str] = []
    for line in req_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue

        # ignore common pip flags; keep it simple & safe for "runtime deps"
        if line.startswith(("-", "--")):
            continue

        dependencies.append(line)

    return dependencies


def _generate_package_files(
    build_dir: Path,
    pkg_dir: Path,
    dependencies: list[str],
) -> None:
    # Generate the __init__.py file in the package directory
    init_path = pkg_dir / "__init__.py"
    if not init_path.exists():
        init_content = "# Auto-generated package initializer\n"
        init_path.write_text(init_content, encoding="utf-8")

    # Generate the pyproject.toml file
    pyproject_path = build_dir / "pyproject.toml"
    deps_str = ",\n  ".join([f'"{d}"' for d in dependencies])
    pyproject_content = f"""
# Auto-generated pyproject.toml file for the module
[project]
name = "voicerun-module"
version = "0.0.0"
description = "Dynamically wrapped module for VoiceRun code"
requires-python = ">=3.12"
dependencies = [
{deps_str}
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/voicerun_module"]
"""
    pyproject_path.write_text(pyproject_content, encoding="utf-8")


async def _install_module(
    source_dir: str | Path,
    build_dir: str | Path,
) -> Path:
    """
    Take a flat user module bundle in source_dir and create an installable package in build_dir.
    Then install it into the current environment using uv.

    Returns the created project root path (build_dir).
    """
    source_dir = Path(source_dir).resolve()
    build_dir = Path(build_dir).resolve()

    # Clean the build directory if it exists
    if build_dir.exists():
        shutil.rmtree(build_dir)

    # Create a new build directory
    build_dir.mkdir(parents=True, exist_ok=True)

    # Create the scaffolding for the module
    pkg_dir = build_dir / "src" / "voicerun_module"
    pkg_dir.mkdir(parents=True, exist_ok=True)

    # Copy only Python files (*.py) from downloaded code to new module directory
    for p in source_dir.glob("*.py"):
        if p.is_file():
            shutil.copy2(p, pkg_dir / p.name)

    # If user supplied pyproject.toml, copy it; otherwise generate one from requirements.txt.
    # If user supplied pyproject.toml, also read its dependencies.
    user_pyproject = source_dir / "pyproject.toml"
    user_requirements = source_dir / "requirements.txt"

    # Get the dependencies from the user's pyproject.toml or requirements.txt in that order.
    dependencies: list[str] = []
    if user_pyproject.exists():
        dependencies = _get_dependencies_from_pyproject(user_pyproject)
    elif user_requirements.exists():
        dependencies = _get_dependencies_from_requirements(user_requirements)

    # Generate the package files
    _generate_package_files(
        build_dir,
        pkg_dir,
        dependencies=dependencies,
    )

    # Ensure a venv exists before installing
    venv_path = _ensure_venv()

    # Install using uv into the venv
    subprocess.run(
        ["uv", "pip", "install", "--python", str(venv_path / "bin" / "python"), str(build_dir)],
        check=True,
        capture_output=not _verbose,
    )

    return build_dir


async def load_module(agent_id: str):
    """
    Load the module from the current directory and install it into the global .voicerun directory.
    Builds are organized by agent ID under ~/.voicerun/builds/<agent_id>/
    """
    current_dir = Path.cwd()

    # Use global builds directory organized by agent ID
    agent_build_dir = BUILDS_DIR / agent_id
    module_dir = agent_build_dir / "module"
    module_dir.mkdir(parents=True, exist_ok=True)

    logger.debug(f"Loading module from local path: {current_dir}")
    logger.debug(f"Installing to global builds directory: {agent_build_dir}")

    await _load_module_from_local_path(current_dir, module_dir)

    # Install the module so that it is importable by the websocket server.
    build_dir = agent_build_dir / "build"
    await _install_module(module_dir, build_dir)

    # If an environment is provided, get the variables for the agent environment.
    environment = os.environ.get("ENVIRONMENT")
    if environment:
        await _get_secrets_for_agent(agent_id, environment)
