"""Module providing task_utils functionality."""

import logging
import os
import shutil
import requests
import zipfile
from typing import Optional
from matrice_common.utils import log_errors
from matrice_compute.scaling import Scaling


def _download_file(url: str, dest_path: str, scaling: Optional[Scaling] = None, max_retries: int = 2) -> None:
    """Download a file from URL using requests with retry on 403.

    Uses requests.get(stream=True) for better error handling and authentication
    support compared to urllib.request.urlretrieve.

    Args:
        url: URL to download from
        dest_path: Local path to save the file
        scaling: Optional Scaling instance for refreshing presigned URLs on 403
        max_retries: Maximum number of retries on 403 (presigned URL expiry)
    """
    # Ensure parent directory exists with proper permissions
    parent_dir = os.path.dirname(dest_path)
    if parent_dir:
        os.makedirs(parent_dir, mode=0o755, exist_ok=True)

    for attempt in range(max_retries + 1):
        try:
            response = requests.get(url, stream=True, timeout=(10, 300))

            if response.status_code == 403 and attempt < max_retries and scaling:
                logging.warning(
                    "Got 403 on download attempt %d/%d, refreshing presigned URL",
                    attempt + 1, max_retries + 1
                )
                refreshed_url = refresh_url_if_needed(url, scaling)
                if refreshed_url and refreshed_url != url:
                    url = refreshed_url
                    continue

            response.raise_for_status()

            with open(dest_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)

            # Fix file permissions
            try:
                os.chmod(dest_path, 0o644)
            except OSError as e:
                logging.debug("Could not set permissions on %s: %s", dest_path, e)

            logging.info("Downloaded %s successfully", dest_path)
            return

        except requests.exceptions.HTTPError as e:
            if response.status_code == 403 and attempt < max_retries:
                continue  # Already handled above
            logging.error("HTTP error downloading %s: %s", url, e)
            raise
        except Exception as e:
            logging.error("Error downloading %s (attempt %d): %s", url, attempt + 1, e)
            if attempt >= max_retries:
                raise


@log_errors(raise_exception=False, log_error=True, default_return=None)
def refresh_url_if_needed(url: Optional[str], scaling: Optional[Scaling] = None) -> Optional[str]:
    """Refresh a presigned URL if it appears to be expired or about to expire.

    This function attempts to refresh presigned URLs for model codebase and requirements
    to ensure they are valid before downloading.

    Args:
        url: The URL to potentially refresh. If None or empty, returns None.
        scaling: The Scaling instance to use for API calls. If None, returns original URL.

    Returns:
        The refreshed URL if successful, or the original URL if refresh fails or is not needed.
    """
    if not url:
        return None

    if not scaling:
        logging.warning("No scaling instance provided, returning original URL")
        return url

    logging.info("Attempting to refresh presigned URL")
    try:
        result = scaling.refresh_presigned_url(url)
        if result is None:
            logging.warning("refresh_presigned_url returned None, using original URL")
            return url
        refreshed_url, error, message = result
        if error:
            logging.warning(f"Failed to refresh presigned URL: {message}. Using original URL.")
            return url
        if refreshed_url:
            logging.info("Successfully refreshed presigned URL")
            return refreshed_url
        else:
            logging.warning("Refresh returned empty URL, using original URL")
            return url
    except Exception as e:
        logging.warning(f"Exception while refreshing presigned URL: {e}. Using original URL.")
        return url


@log_errors(raise_exception=True, log_error=True)
def setup_workspace_and_run_task(
    work_fs: str,
    action_id: str,
    model_codebase_url: str,
    model_codebase_requirements_url: Optional[str] = None,
    scaling: Optional[Scaling] = None,
) -> None:
    """Set up workspace and run task with provided parameters.

    Args:
        work_fs (str): Working filesystem path.
        action_id (str): Unique identifier for the action.
        model_codebase_url (str): URL to download model codebase from.
        model_codebase_requirements_url (Optional[str]): URL to download requirements from. Defaults to None.
        scaling (Optional[Scaling]): Scaling instance for refreshing presigned URLs. Defaults to None.

    Returns:
        None
    """
    workspace_dir = f"{work_fs}/{action_id}"
    codebase_zip_path = f"{workspace_dir}/file.zip"
    requirements_txt_path = f"{workspace_dir}/requirements.txt"

    # Clean up existing workspace to ensure fresh codebase on re-runs
    if os.path.exists(workspace_dir):
        logging.info("Cleaning up existing workspace directory: %s", workspace_dir)
        shutil.rmtree(workspace_dir)

    os.makedirs(workspace_dir, mode=0o755, exist_ok=True)

    # Refresh presigned URLs before downloading to ensure they are valid
    refreshed_codebase_url = refresh_url_if_needed(model_codebase_url, scaling)
    if refreshed_codebase_url:
        model_codebase_url = refreshed_codebase_url

    # Download codebase ZIP file using requests with retry
    _download_file(model_codebase_url, codebase_zip_path, scaling=scaling)

    # Extract ZIP file with overwrite
    with zipfile.ZipFile(codebase_zip_path, 'r') as zip_ref:
        for member in zip_ref.namelist():
            member_path = os.path.realpath(os.path.join(workspace_dir, member))
            if not member_path.startswith(os.path.realpath(workspace_dir) + os.sep):
                raise ValueError(f"Attempted path traversal in zip: {member}")
        zip_ref.extractall(workspace_dir)

    # If the ZIP contained a single top-level directory, move its contents
    # up one level so deploy.py / train.py etc. sit directly in workspace_dir.
    # IMPORTANT: use shutil.move (not flat-walk) to preserve nested package
    # directories like inference/, which are required for Python imports.
    entries = [
        e for e in os.listdir(workspace_dir)
        if e != os.path.basename(codebase_zip_path)
    ]
    if len(entries) == 1:
        top_dir = os.path.join(workspace_dir, entries[0])
        if os.path.isdir(top_dir):
            logging.info(
                "Stripping single top-level directory '%s' from extracted codebase",
                entries[0],
            )
            for item in os.listdir(top_dir):
                src = os.path.join(top_dir, item)
                dst = os.path.join(workspace_dir, item)
                if os.path.exists(dst):
                    if os.path.isdir(dst):
                        shutil.rmtree(dst)
                    else:
                        try:
                            if not os.access(dst, os.W_OK):
                                os.chmod(dst, 0o644)
                        except OSError:
                            pass
                        os.remove(dst)
                try:
                    shutil.move(src, dst)
                except PermissionError:
                    logging.warning("Permission denied moving %s -> %s, attempting chmod fix", src, dst)
                    try:
                        if os.path.isdir(src):
                            os.chmod(src, 0o755)
                        else:
                            os.chmod(src, 0o644)
                        os.chmod(os.path.dirname(dst), 0o755)
                        shutil.move(src, dst)
                    except PermissionError:
                        logging.error("Cannot fix permissions for %s -> %s, may need root access", src, dst)
                        raise
            # Remove the now-empty top-level directory
            try:
                os.rmdir(top_dir)
            except OSError:
                pass

    # Download requirements.txt if URL is provided
    if model_codebase_requirements_url:
        # Refresh presigned URL for requirements before downloading
        refreshed_requirements_url = refresh_url_if_needed(model_codebase_requirements_url, scaling)
        if refreshed_requirements_url:
            model_codebase_requirements_url = refreshed_requirements_url
        _download_file(model_codebase_requirements_url, requirements_txt_path, scaling=scaling)
