"""
HTTP client for Modal worker communication.

Handles:
- Tool execution requests via async submit + poll pattern
- API key authentication via Authorization header
- Job ID generation and tracking
- Response parsing and error handling
- Exponential backoff with jitter for polling

Note: Job/artifact tracking is handled server-side by Modal workers.
"""

import asyncio
import httpx
import random
import string
import time
from typing import Any, Optional
from uuid import uuid4

from amina_cli.auth import get_api_key, AuthError
from amina_cli.registry import get_submit_endpoint, get_status_endpoint, get_queued_status_endpoint, ToolNotFoundError


# Default polling configuration
DEFAULT_POLL_INTERVAL = 3.0  # Initial poll interval in seconds
DEFAULT_MAX_POLL_INTERVAL = 15.0  # Max poll interval (with backoff)
DEFAULT_BACKOFF_MULTIPLIER = 1.5  # Exponential backoff multiplier
DEFAULT_JITTER = 0.1  # Jitter range (+-10%)
DEFAULT_TIMEOUT = 36000.0  # 10 hours max for long-running tools
DEFAULT_MAX_NETWORK_RETRIES = 5  # Max retries for transient network errors


class ClientError(Exception):
    """Base exception for client errors."""

    pass


class AuthenticationError(ClientError):
    """Raised when API key is invalid or missing."""

    pass


class ToolExecutionError(ClientError):
    """Raised when tool execution fails."""

    pass


class InsufficientCreditsError(ClientError):
    """Raised when user has insufficient credits for the tool."""

    def __init__(self, required: float, remaining: float, upgrade_url: str):
        self.required = required
        self.remaining = remaining
        self.upgrade_url = upgrade_url
        super().__init__(
            f"Insufficient credits. This tool requires {required} credits, "
            f"but you have {remaining:.1f} remaining. "
            f"Upgrade at: {upgrade_url}"
        )


class JobLimitError(ClientError):
    """Raised when user has reached job concurrency limits."""

    def __init__(self, limit_type: str, current: int, limit: int, message: str):
        self.limit_type = limit_type
        self.current = current
        self.limit = limit
        super().__init__(message)


def generate_job_id() -> str:
    """Generate a unique job ID for tracking."""
    return str(uuid4())


def generate_job_name(tool_name: str) -> str:
    """Generate a 4-letter random job name."""
    return "".join(random.choices(string.ascii_uppercase, k=4))


def _build_payload(tool_name: str, params: dict[str, Any], job_id: str, job_name: str) -> dict[str, Any]:
    """
    Build the request payload in the format expected by the CLI Gateway.

    All tools use the same nested format:
    - tool_input: Tool-specific parameters including job_id, job_name
    - storage_config: Storage configuration for input/output files

    Args:
        tool_name: Name of the tool
        params: Tool-specific parameters from CLI
        job_id: Generated job ID
        job_name: Generated job name

    Returns:
        Payload dict ready for the gateway API
    """
    return {
        "tool_input": {
            "job_id": job_id,
            "job_name": job_name,
            **params,
        },
        "storage_config": {
            "input_bucket": "user-files",
            "output_bucket": "user-files",
            "output_dir": f"cli/{job_id}/",
        },
    }


def _calculate_poll_interval(
    base_interval: float,
    attempt: int,
    multiplier: float,
    max_interval: float,
    jitter: float,
) -> float:
    """
    Calculate poll interval with exponential backoff and jitter.

    Args:
        base_interval: Initial poll interval
        attempt: Current attempt number (0-indexed)
        multiplier: Backoff multiplier (e.g., 1.5)
        max_interval: Maximum poll interval
        jitter: Jitter range (e.g., 0.1 for +-10%)

    Returns:
        Poll interval with backoff and jitter applied
    """
    # Exponential backoff: base * multiplier^attempt
    interval = base_interval * (multiplier**attempt)

    # Cap at max interval
    interval = min(interval, max_interval)

    # Apply jitter (+-jitter%)
    jitter_range = interval * jitter
    interval += random.uniform(-jitter_range, jitter_range)

    return interval


# Callback type for queue position updates
QueuePositionCallback = Any  # Callable[[int], None] but avoiding import


async def _poll_queued_until_submitted(
    client: httpx.AsyncClient,
    job_id: str,
    timeout: float,
    start_time: float,
    poll_interval: float = 5.0,  # Longer interval for queue polling
    max_poll_interval: float = 30.0,
    on_position_update: Optional[QueuePositionCallback] = None,
) -> tuple[str, str]:
    """
    Poll queued_job_status until job is submitted.

    Args:
        client: HTTP client
        job_id: Job ID to poll for
        timeout: Total timeout in seconds
        start_time: When the job was submitted
        poll_interval: Initial poll interval
        max_poll_interval: Maximum poll interval
        on_position_update: Optional callback when queue position changes

    Returns:
        Tuple of (call_id, user_id) when job is submitted

    Raises:
        ToolExecutionError: If job fails or times out in queue
    """
    queued_endpoint = get_queued_status_endpoint()
    poll_attempt = 0
    last_position = -1

    while True:
        # Check timeout
        elapsed = time.time() - start_time
        if elapsed > timeout:
            raise ToolExecutionError(f"Job timed out after {int(elapsed)}s waiting in queue. Job ID: {job_id}")

        # Calculate poll interval (slower backoff for queue polling)
        interval = _calculate_poll_interval(
            poll_interval,
            poll_attempt,
            1.3,  # Slower backoff for queue
            max_poll_interval,
            DEFAULT_JITTER,
        )
        await asyncio.sleep(interval)
        poll_attempt += 1

        # Poll queued status
        status_url = f"{queued_endpoint}?job_id={job_id}"

        try:
            response = await client.get(status_url)
        except httpx.RequestError as e:
            # Log transient error and continue polling
            print(f"Queue poll network error (will retry): {type(e).__name__}: {e}")
            continue

        if response.status_code != 200:
            # Unexpected error
            raise ToolExecutionError(f"Queue status check failed: HTTP {response.status_code}")

        result = response.json()
        status = result.get("status")

        if status == "queued":
            # Still in queue - update position if callback provided
            position = result.get("queue_position", 0)
            if on_position_update and position != last_position:
                on_position_update(position)
                last_position = position
            continue

        elif status == "submitted":
            # Job has been spawned - return call_id
            call_id = result.get("call_id")
            if not call_id:
                raise ToolExecutionError("Job submitted but no call_id returned")
            return call_id, ""  # user_id not in this response

        elif status == "running" or status == "completed":
            # Job already running/done - return call_id
            call_id = result.get("call_id")
            if not call_id:
                raise ToolExecutionError(f"Job {status} but no call_id returned")
            return call_id, ""

        elif status == "failed":
            # Job failed (possibly expired)
            error = result.get("error", "Job failed in queue")
            raise ToolExecutionError(f"Job failed: {error}")

        elif status == "not_found":
            raise ToolExecutionError(f"Job not found: {job_id}")

        else:
            # Unknown status - log and continue polling
            print(f"Unknown queue status '{status}' for job {job_id}, will retry")
            continue


async def run_tool(
    tool_name: str,
    params: dict[str, Any],
    api_key: Optional[str] = None,
    poll_interval: float = DEFAULT_POLL_INTERVAL,
    max_poll_interval: float = DEFAULT_MAX_POLL_INTERVAL,
    timeout: float = DEFAULT_TIMEOUT,
    max_network_retries: int = DEFAULT_MAX_NETWORK_RETRIES,
    on_queue_position_update: Optional[QueuePositionCallback] = None,
) -> dict[str, Any]:
    """
    Execute a tool via async submit + poll pattern.

    This function submits a job to the gateway, then polls for completion.
    The pattern avoids Modal's 150-second HTTP timeout for long-running tools.

    If the concurrent job limit is reached, the job will be queued and this
    function will poll until the queue processor spawns it.

    Args:
        tool_name: Name of the tool to run (e.g., "esmfold", "pdb-cleaner")
        params: Tool-specific parameters
        api_key: Optional API key (uses stored key if not provided)
        poll_interval: Initial poll interval in seconds (default 3s)
        max_poll_interval: Maximum poll interval with backoff (default 15s)
        timeout: Maximum total time to wait (default 10 hours)
        max_network_retries: Max retries for transient network errors (default 5)
        on_queue_position_update: Optional callback when queue position changes

    Returns:
        Tool execution result with status, output_files, etc.

    Raises:
        AuthenticationError: If API key is missing or invalid
        ToolNotFoundError: If tool is not found in registry
        InsufficientCreditsError: If user has insufficient credits
        ToolExecutionError: If tool execution fails
    """
    # Get API key
    if api_key is None:
        try:
            api_key = get_api_key()
        except AuthError as e:
            raise AuthenticationError(str(e))

    # Get endpoints
    try:
        submit_endpoint = get_submit_endpoint(tool_name)
        status_endpoint = get_status_endpoint()
    except ToolNotFoundError:
        raise

    # Build request payload
    job_id = generate_job_id()
    job_name = generate_job_name(tool_name)
    payload = _build_payload(tool_name, params, job_id, job_name)

    start_time = time.time()

    async with httpx.AsyncClient(timeout=60.0) as client:
        # ═══════════════════════════════════════════════════════════════════════
        # PHASE 1: SUBMIT JOB
        # ═══════════════════════════════════════════════════════════════════════
        try:
            response = await client.post(
                submit_endpoint,
                json=payload,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
            )
        except httpx.RequestError as e:
            raise ToolExecutionError(f"Network error submitting job: {str(e)}")

        # Handle submit errors
        if response.status_code == 401:
            raise AuthenticationError("Invalid API key. Get a new one at: https://app.aminoanalytica.com/settings/api")
        if response.status_code == 403:
            raise AuthenticationError("API key has been revoked. Generate a new one in settings.")

        if response.status_code == 402:
            try:
                detail = response.json().get("detail", {})
                raise InsufficientCreditsError(
                    required=detail.get("required", 0),
                    remaining=detail.get("remaining", 0),
                    upgrade_url=detail.get("upgrade_url", "https://app.aminoanalytica.com/settings/billing"),
                )
            except (ValueError, KeyError):
                raise InsufficientCreditsError(
                    required=0,
                    remaining=0,
                    upgrade_url="https://app.aminoanalytica.com/settings/billing",
                )

        if response.status_code == 422:
            try:
                detail = response.json().get("detail", {})
                errors = detail.get("errors", [])
                raise ToolExecutionError(f"Input validation failed: {', '.join(errors)}")
            except (ValueError, KeyError):
                raise ToolExecutionError(f"Input validation failed: {response.text}")

        if response.status_code == 429:
            try:
                detail = response.json().get("detail", {})
                raise JobLimitError(
                    limit_type=detail.get("limit_type", "unknown"),
                    current=detail.get("current", 0),
                    limit=detail.get("limit", 0),
                    message=detail.get("message", "Job limit reached. Please wait for running jobs to complete."),
                )
            except (ValueError, KeyError):
                raise JobLimitError(
                    limit_type="unknown",
                    current=0,
                    limit=0,
                    message="Job limit reached. Please wait for running jobs to complete.",
                )

        if response.status_code == 404:
            raise ToolNotFoundError(f"Tool '{tool_name}' not found")

        if response.status_code == 503:
            raise ToolExecutionError(f"Tool '{tool_name}' is not currently deployed")

        if response.status_code != 200:
            raise ToolExecutionError(f"Submit failed with HTTP {response.status_code}: {response.text}")

        # Extract response info
        submit_result = response.json()
        status = submit_result.get("status")
        user_id = submit_result.get("user_id", "")

        # ═══════════════════════════════════════════════════════════════════════
        # PHASE 1.5: HANDLE QUEUED JOBS
        # ═══════════════════════════════════════════════════════════════════════
        if status == "queued":
            # Job was queued due to concurrent limit - poll until submitted
            queue_position = submit_result.get("queue_position", 0)
            if on_queue_position_update:
                on_queue_position_update(queue_position)

            call_id, _ = await _poll_queued_until_submitted(
                client,
                job_id,
                timeout,
                start_time,
                poll_interval=5.0,  # Slower for queue
                max_poll_interval=30.0,
                on_position_update=on_queue_position_update,
            )
        else:
            # Job was submitted immediately
            call_id = submit_result.get("call_id")
            if not call_id:
                raise ToolExecutionError("Submit succeeded but no call_id returned")

        # ═══════════════════════════════════════════════════════════════════════
        # PHASE 2: POLL FOR COMPLETION
        # ═══════════════════════════════════════════════════════════════════════
        poll_attempt = 0
        consecutive_network_errors = 0

        while True:
            # Check timeout
            elapsed = time.time() - start_time
            if elapsed > timeout:
                raise ToolExecutionError(
                    f"Tool '{tool_name}' timed out after {int(elapsed)}s. "
                    f"The job may still be running in the background."
                )

            # Calculate poll interval with backoff
            interval = _calculate_poll_interval(
                poll_interval,
                poll_attempt,
                DEFAULT_BACKOFF_MULTIPLIER,
                max_poll_interval,
                DEFAULT_JITTER,
            )
            await asyncio.sleep(interval)
            poll_attempt += 1

            # Poll for status
            status_url = f"{status_endpoint}?call_id={call_id}&job_id={job_id}&user_id={user_id}&tool_name={tool_name}"

            try:
                response = await client.get(status_url)
                consecutive_network_errors = 0  # Reset on success
            except httpx.RequestError as e:
                consecutive_network_errors += 1
                if consecutive_network_errors >= max_network_retries:
                    raise ToolExecutionError(
                        f"Network error polling job status: {str(e)}. The job may still be running. Job ID: {job_id}"
                    )
                # Transient error - continue polling
                continue

            # Handle poll response
            if response.status_code == 202:
                # Job still running - continue polling
                continue

            if response.status_code == 200:
                # Job completed
                result = response.json()
                if result.get("status") == "completed":
                    tool_result = result.get("result", {})
                    # Add job tracking info
                    tool_result["job_id"] = job_id
                    tool_result["job_name"] = job_name
                    return tool_result
                # Unexpected status
                raise ToolExecutionError(f"Unexpected poll response: {result}")

            if response.status_code == 404:
                try:
                    error = response.json().get("error", "Job not found")
                except ValueError:
                    error = "Job not found"
                raise ToolExecutionError(f"Job not found: {error}")

            if response.status_code == 410:
                try:
                    error = response.json().get("error", "Job results have expired")
                except ValueError:
                    error = "Job results have expired"
                raise ToolExecutionError(f"Job results expired: {error}")

            if response.status_code == 500:
                try:
                    error = response.json().get("error", "Job failed")
                except ValueError:
                    error = response.text
                raise ToolExecutionError(f"Job failed: {error}")

            # Unexpected status code
            raise ToolExecutionError(f"Unexpected poll status {response.status_code}: {response.text}")


def run_tool_sync(
    tool_name: str,
    params: dict[str, Any],
    api_key: Optional[str] = None,
    timeout: float = DEFAULT_TIMEOUT,
    on_queue_position_update: Optional[QueuePositionCallback] = None,
) -> dict[str, Any]:
    """
    Synchronous wrapper for run_tool.

    Use this from non-async CLI commands.

    Args:
        tool_name: Name of the tool to run
        params: Tool-specific parameters
        api_key: Optional API key
        timeout: Maximum total time to wait in seconds
        on_queue_position_update: Optional callback when queue position changes

    Returns:
        Tool execution result
    """
    return asyncio.run(
        run_tool(
            tool_name,
            params,
            api_key,
            timeout=timeout,
            on_queue_position_update=on_queue_position_update,
        )
    )


# ═══════════════════════════════════════════════════════════════════════════════
# BACKGROUND JOB FUNCTIONS
# ═══════════════════════════════════════════════════════════════════════════════


async def submit_tool(
    tool_name: str,
    params: dict[str, Any],
    api_key: Optional[str] = None,
) -> dict[str, Any]:
    """
    Submit a tool job and return immediately without waiting for completion.

    This is the first phase of the async submit + poll pattern. Use this for
    background job submission where you want to poll for status separately.

    Args:
        tool_name: Name of the tool to run (e.g., "esmfold", "pdb-cleaner")
        params: Tool-specific parameters
        api_key: Optional API key (uses stored key if not provided)

    Returns:
        Dict with job tracking info: {job_id, call_id, user_id, job_name}

    Raises:
        AuthenticationError: If API key is missing or invalid
        ToolNotFoundError: If tool is not found in registry
        InsufficientCreditsError: If user has insufficient credits
        ToolExecutionError: If submission fails
    """
    # Get API key
    if api_key is None:
        try:
            api_key = get_api_key()
        except AuthError as e:
            raise AuthenticationError(str(e))

    # Get endpoint
    try:
        submit_endpoint = get_submit_endpoint(tool_name)
    except ToolNotFoundError:
        raise

    # Build request payload
    job_id = generate_job_id()
    job_name = generate_job_name(tool_name)
    payload = _build_payload(tool_name, params, job_id, job_name)

    async with httpx.AsyncClient(timeout=60.0) as client:
        try:
            response = await client.post(
                submit_endpoint,
                json=payload,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
            )
        except httpx.RequestError as e:
            raise ToolExecutionError(f"Network error submitting job: {str(e)}")

        # Handle submit errors
        if response.status_code == 401:
            raise AuthenticationError("Invalid API key. Get a new one at: https://app.aminoanalytica.com/settings/api")
        if response.status_code == 403:
            raise AuthenticationError("API key has been revoked. Generate a new one in settings.")

        if response.status_code == 402:
            try:
                detail = response.json().get("detail", {})
                raise InsufficientCreditsError(
                    required=detail.get("required", 0),
                    remaining=detail.get("remaining", 0),
                    upgrade_url=detail.get("upgrade_url", "https://app.aminoanalytica.com/settings/billing"),
                )
            except (ValueError, KeyError):
                raise InsufficientCreditsError(
                    required=0,
                    remaining=0,
                    upgrade_url="https://app.aminoanalytica.com/settings/billing",
                )

        if response.status_code == 422:
            try:
                detail = response.json().get("detail", {})
                errors = detail.get("errors", [])
                raise ToolExecutionError(f"Input validation failed: {', '.join(errors)}")
            except (ValueError, KeyError):
                raise ToolExecutionError(f"Input validation failed: {response.text}")

        if response.status_code == 429:
            try:
                detail = response.json().get("detail", {})
                raise JobLimitError(
                    limit_type=detail.get("limit_type", "unknown"),
                    current=detail.get("current", 0),
                    limit=detail.get("limit", 0),
                    message=detail.get("message", "Job limit reached. Please wait for running jobs to complete."),
                )
            except (ValueError, KeyError):
                raise JobLimitError(
                    limit_type="unknown",
                    current=0,
                    limit=0,
                    message="Job limit reached. Please wait for running jobs to complete.",
                )

        if response.status_code == 404:
            raise ToolNotFoundError(f"Tool '{tool_name}' not found")

        if response.status_code == 503:
            raise ToolExecutionError(f"Tool '{tool_name}' is not currently deployed")

        if response.status_code != 200:
            raise ToolExecutionError(f"Submit failed with HTTP {response.status_code}: {response.text}")

        # Extract tracking info
        submit_result = response.json()
        status = submit_result.get("status", "submitted")
        call_id = submit_result.get("call_id")

        # Handle both submitted and queued responses
        # Treat as queued if status is "queued" OR if no call_id was returned
        if status == "queued" or not call_id:
            return {
                "job_id": job_id,
                "job_name": job_name,
                "status": "queued",
                "queue_position": submit_result.get("queue_position", 0),
                "user_id": submit_result.get("user_id", ""),
                "tool_name": tool_name,
            }
        else:
            return {
                "job_id": job_id,
                "job_name": job_name,
                "status": "submitted",
                "call_id": call_id,
                "user_id": submit_result.get("user_id", ""),
                "tool_name": tool_name,
            }


def submit_tool_sync(
    tool_name: str,
    params: dict[str, Any],
    api_key: Optional[str] = None,
) -> dict[str, Any]:
    """
    Synchronous wrapper for submit_tool.

    Args:
        tool_name: Name of the tool to run
        params: Tool-specific parameters
        api_key: Optional API key

    Returns:
        Dict with job tracking info: {job_id, call_id, user_id, job_name, tool_name}
    """
    return asyncio.run(submit_tool(tool_name, params, api_key))


async def check_job_status(
    call_id: str,
    job_id: str,
    user_id: str = "",
    tool_name: str = "",
) -> dict[str, Any]:
    """
    Check the status of a submitted job (single poll, no retry loop).

    Args:
        call_id: Modal function call ID from submit_tool response
        job_id: Job ID from submit_tool response
        user_id: User ID from submit_tool response (optional)
        tool_name: Tool name for error messages (optional)

    Returns:
        Dict with status info:
        - {"status": "running"} if job is still running (HTTP 202)
        - {"status": "completed", "result": {...}} if job finished successfully
        - {"status": "failed", "error": "..."} if job failed

    Raises:
        ToolExecutionError: If status check fails unexpectedly
    """
    status_endpoint = get_status_endpoint()
    status_url = f"{status_endpoint}?call_id={call_id}&job_id={job_id}&user_id={user_id}&tool_name={tool_name}"

    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            response = await client.get(status_url)
        except httpx.RequestError as e:
            raise ToolExecutionError(f"Network error checking job status: {str(e)}")

        if response.status_code == 202:
            # Job still running
            return {"status": "running"}

        if response.status_code == 200:
            # Job completed
            result = response.json()
            if result.get("status") == "completed":
                tool_result = result.get("result", {})
                tool_result["job_id"] = job_id
                return {"status": "completed", "result": tool_result}
            # Unexpected status in 200 response
            return {"status": "unknown", "raw": result}

        if response.status_code == 404:
            try:
                error = response.json().get("error", "Job not found")
            except ValueError:
                error = "Job not found"
            return {"status": "failed", "error": error}

        if response.status_code == 410:
            try:
                error = response.json().get("error", "Job results have expired")
            except ValueError:
                error = "Job results have expired"
            return {"status": "failed", "error": error}

        if response.status_code == 500:
            try:
                error = response.json().get("error", "Job failed")
            except ValueError:
                error = response.text
            return {"status": "failed", "error": error}

        # Unexpected status code
        return {"status": "failed", "error": f"Unexpected status {response.status_code}: {response.text}"}


def check_job_status_sync(
    call_id: str,
    job_id: str,
    user_id: str = "",
    tool_name: str = "",
) -> dict[str, Any]:
    """
    Synchronous wrapper for check_job_status.

    Args:
        call_id: Modal function call ID
        job_id: Job ID
        user_id: User ID (optional)
        tool_name: Tool name (optional)

    Returns:
        Dict with status info
    """
    return asyncio.run(check_job_status(call_id, job_id, user_id, tool_name))


async def check_queued_job_status(job_id: str) -> dict[str, Any]:
    """
    Check the status of a queued job.

    Args:
        job_id: Job ID from submit_tool response

    Returns:
        Dict with status info:
        - {"status": "queued", "queue_position": N} if still waiting
        - {"status": "submitted", "call_id": "..."} if spawned
        - {"status": "failed", "error": "..."} if failed/expired
    """
    queued_endpoint = get_queued_status_endpoint()
    status_url = f"{queued_endpoint}?job_id={job_id}"

    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            response = await client.get(status_url)
        except httpx.RequestError as e:
            raise ToolExecutionError(f"Network error checking queued job status: {str(e)}")

        if response.status_code != 200:
            return {"status": "error", "error": f"HTTP {response.status_code}"}

        return response.json()


def check_queued_job_status_sync(job_id: str) -> dict[str, Any]:
    """
    Synchronous wrapper for check_queued_job_status.

    Args:
        job_id: Job ID

    Returns:
        Dict with status info
    """
    return asyncio.run(check_queued_job_status(job_id))


def check_endpoint_health(tool_name: str) -> bool:
    """
    Check if a tool's endpoint is reachable.

    Args:
        tool_name: Name of the tool to check

    Returns:
        True if endpoint is reachable
    """
    try:
        from amina_cli.registry import _get_gateway_base

        base = _get_gateway_base()
        health_url = f"{base}-health.modal.run"

        with httpx.Client(timeout=10.0) as client:
            response = client.get(health_url)
            return response.status_code < 500
    except Exception:
        return False
