"""Lambda handler shim for AlphaAgent code-interpreter sessions.

This module is the ENTRYPOINT for every environment Lambda function.
It receives a JSON event from the LambdaBackend and executes Python code
or shell commands inside the Lambda execution environment.

Scripts and output files are written to /tmp (always writable in Lambda).
After execution, all generated files are uploaded to S3 via direct API calls.

Connector credentials are fetched from Secrets Manager at invocation time
and injected as environment variables before running user code.
"""

import json
import os
import subprocess
import sys
import time
import traceback
from typing import Any, Dict, List, Tuple

import boto3

MAX_OUTPUT_BYTES = 1024 * 1024  # 1 MB cap on stdout/stderr returned to caller

# Checkpoint/continuation contract (see aa_env.py). When a script checkpoints and
# asks to be re-invoked it exits with this code; the coder harness re-invokes a
# fresh Lambda with resume=true. Kept in one place so the three parties (this
# handler, aa_env, coder_worker) agree.
EXIT_CONTINUE = 75
# Seconds reserved AFTER the script exits for the workdir→S3 upload, and seconds
# of grace before the hard kill in which a cooperative script should checkpoint.
# Overridable via env (folded from common-config into the Lambda by the backend).
_DEFAULT_UPLOAD_RESERVE_S = 60
_DEFAULT_YIELD_GRACE_S = 30


def _load_connector_credentials(agent_id: str, connector_ids: list) -> Dict[str, str]:
    """Fetch connector credentials from Secrets Manager and return as env vars."""
    if not connector_ids:
        return {}

    try:
        sm = boto3.client("secretsmanager")
        secret_name = f"alphaagent-connector-creds-{agent_id}"
        print(f"[HANDLER] Fetching connector secret: {secret_name}")
        resp = sm.get_secret_value(SecretId=secret_name)
        creds = json.loads(resp["SecretString"])
        env_vars: Dict[str, str] = {}
        for key, value in creds.items():
            env_vars[key] = str(value) if value is not None else ""
        print(f"[HANDLER] Loaded {len(env_vars)} connector creds")
        return env_vars
    except Exception as exc:
        print(f"[HANDLER] ERROR loading connector creds: {exc}")
        return {}


def _cap_output(data: bytes) -> str:
    """Decode and truncate output to stay within Lambda response limits."""
    text = data.decode("utf-8", errors="replace")
    if len(text) > MAX_OUTPUT_BYTES:
        return text[:MAX_OUTPUT_BYTES] + "\n[OUTPUT TRUNCATED]"
    return text


def _s3_client():
    return boto3.client("s3")


def _upload_workdir_files(bucket: str, prefix: str, workdir: str) -> List[str]:
    """Upload all files from workdir to S3. Returns list of uploaded relative paths."""
    s3 = _s3_client()
    uploaded = []

    for dirpath, dirnames, filenames in os.walk(workdir):
        # Skip hidden directories
        dirnames[:] = [d for d in dirnames if not d.startswith(".")]
        for filename in filenames:
            if filename.startswith("."):
                continue
            local_path = os.path.join(dirpath, filename)
            rel_path = os.path.relpath(local_path, workdir)
            s3_key = f"{prefix.rstrip('/')}/{rel_path}"
            try:
                with open(local_path, "rb") as f:
                    content = f.read()
                s3.put_object(Bucket=bucket, Key=s3_key, Body=content)
                uploaded.append(rel_path)
            except Exception as e:
                print(f"[HANDLER] WARNING: failed to upload {rel_path}: {e}")

    if uploaded:
        print(f"[HANDLER] Uploaded {len(uploaded)} files to s3://{bucket}/{prefix.rstrip('/')}/")
    return uploaded


def _sync_prefix_down(bucket: str, prefix: str, workdir: str) -> int:
    """Download the ENTIRE session prefix into the workdir. Used on checkpoint
    resume so a fresh Lambda's /tmp is reconstituted exactly (manifest + shards +
    prior outputs) — the durable S3 workspace is the source of truth, never /tmp.
    Best-effort: a failure here degrades to recomputing from the last checkpoint.
    """
    if not bucket or not prefix:
        return 0
    s3 = _s3_client()
    count = 0
    try:
        paginator = s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for obj in page.get("Contents", []):
                key = obj["Key"]
                rel = key[len(prefix):].lstrip("/")
                if not rel or rel.endswith("/"):
                    continue
                dst = os.path.join(workdir, rel)
                os.makedirs(os.path.dirname(dst) or workdir, exist_ok=True)
                s3.download_file(bucket, key, dst)
                count += 1
        print(f"[HANDLER] resume: synced {count} file(s) down from s3://{bucket}/{prefix}")
    except Exception as e:
        print(f"[HANDLER] WARNING: resume sync-down failed: {e}")
    return count


def _int_env(name: str, default: int) -> int:
    try:
        return int(os.environ.get(name) or default)
    except (TypeError, ValueError):
        return default


def _compute_budget(event: Dict[str, Any], context: Any) -> Tuple[int, float]:
    """Return ``(hard_timeout_seconds, deadline_epoch)`` for an exec action.

    The hard timeout is the coreutils cap on the script; the deadline is the
    earlier soft point (exported as ``AA_DEADLINE_EPOCH``) at which a cooperative
    checkpointing script should yield, leaving grace to flush + upload.

    Derived from the LIVE Lambda remaining time so a checkpointed chunk is as
    large as safely possible, minus an upload reserve. ``event['timeout']`` (the
    coder's ``CODER_ENV_INVOKE_TIMEOUT_SECONDS`` ceiling) caps it further. Falls
    back to the event timeout when no Lambda context is available (local/tests).
    """
    def _param(key: str, env_name: str, default: int) -> int:
        v = event.get(key)
        if isinstance(v, (int, float)) and v > 0:
            return int(v)
        return _int_env(env_name, default)

    reserve = _param("upload_reserve_s", "AA_UPLOAD_RESERVE_S", _DEFAULT_UPLOAD_RESERVE_S)
    grace = _param("yield_grace_s", "AA_YIELD_GRACE_S", _DEFAULT_YIELD_GRACE_S)
    requested = event.get("timeout")
    remaining = None
    try:
        if context is not None and hasattr(context, "get_remaining_time_in_millis"):
            remaining = context.get_remaining_time_in_millis() / 1000.0
    except Exception:
        remaining = None
    if remaining is not None:
        hard = remaining - reserve
        if isinstance(requested, (int, float)) and requested > 0:
            hard = min(hard, float(requested))
    elif isinstance(requested, (int, float)) and requested > 0:
        hard = float(requested)
    else:
        hard = 300.0
    hard_int = int(max(5, hard))
    deadline = time.time() + max(1, hard_int - grace)
    return hard_int, deadline


def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    """Lambda entrypoint.

    Event schema::

        {
            "action": "exec_python" | "exec_shell" | "write_file" | "read_file" | "list_files",
            "s3_bucket": "alphaagent-workspaces-...",
            "s3_prefix": "workspaces/sessions/{session_id}/workspace/",
            "timeout": 300,           # seconds (for exec actions)
            "code": "...",            # Python source (exec_python)
            "script_name": "...",     # Script filename (exec_python)
            "command": "...",         # Shell command string (exec_shell)
            "path": "...",            # Relative path (file operations)
            "content": "...",         # File content (write_file)
            "agent_id": "...",        # For credential lookup
            "connector_ids": [...]    # Connectors to inject
        }
    """
    action = event.get("action", "")
    s3_bucket = event.get("s3_bucket", "")
    s3_prefix = event.get("s3_prefix", "")
    timeout = event.get("timeout", 300)

    print(f"[HANDLER] entry: action={action}, bucket={s3_bucket}, prefix={s3_prefix}, "
          f"timeout={timeout}, resume={bool(event.get('resume'))}")
    print(f"[HANDLER] uid={os.getuid()}, gid={os.getgid()}")

    if action in ("exec_python", "exec_shell"):
        # Derive the real per-invocation budget from the LIVE Lambda clock and
        # publish the soft deadline so a cooperative checkpointing script (aa_env)
        # can yield before the hard kill. A non-checkpointed script simply runs
        # under the hard timeout as before.
        timeout, deadline = _compute_budget(event, context)
        os.environ["AA_DEADLINE_EPOCH"] = str(deadline)
        print(f"[HANDLER] budget: hard_timeout={timeout}s, deadline_epoch={deadline:.0f}")
        connector_env = _load_connector_credentials(
            event.get("agent_id", ""),
            event.get("connector_ids", []),
        )
        for k, v in connector_env.items():
            os.environ[k] = v
        print(f"[HANDLER] Injected {len(connector_env)} connector env vars")

    try:
        if action == "exec_python":
            return _exec_python(event, s3_bucket, s3_prefix, timeout)
        elif action == "exec_shell":
            return _exec_shell(event, s3_bucket, s3_prefix, timeout)
        elif action == "write_file":
            return _write_file(event, s3_bucket, s3_prefix)
        elif action == "read_file":
            return _read_file(event, s3_bucket, s3_prefix)
        elif action == "list_files":
            return _list_files(event, s3_bucket, s3_prefix)
        else:
            return {"exit_code": 1, "stdout": "", "stderr": f"Unknown action: {action}"}
    except Exception as exc:
        print(f"[HANDLER] Unhandled exception: {traceback.format_exc()}")
        return {
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Handler error: {traceback.format_exc()}",
        }


def _prefetch_files(prefetch: List[str], s3_bucket: str, s3_prefix: str, workdir: str) -> None:
    """Download declared inputs into the workdir before running code.

    Each entry is either a full S3 key or a path relative to ``s3_prefix``. The
    env Lambda's /tmp is ephemeral, so a multi-invoke coder run relies on this to
    make prior artifacts (and conversation-store files) available to the code.
    Best-effort: a missing/failed prefetch never aborts the run.
    """
    if not prefetch or not s3_bucket:
        return
    s3 = _s3_client()
    for entry in prefetch:
        try:
            # Full key (workspaces/.. or conversations/..) is used as-is;
            # anything else is treated as relative to this session's prefix.
            if entry.startswith("workspaces/") or entry.startswith("conversations/"):
                key = entry
            else:
                key = f"{s3_prefix.rstrip('/')}/{entry.lstrip('/')}"
            dst = os.path.join(workdir, os.path.basename(key))
            s3.download_file(s3_bucket, key, dst)
            print(f"[HANDLER] prefetched {key} -> {dst}")
        except Exception as e:
            print(f"[HANDLER] WARNING: prefetch failed for {entry}: {e}")


def _get_session_workdir(s3_prefix: str) -> str:
    """Derive a stable /tmp working directory path from the S3 prefix."""
    # s3_prefix = "workspaces/sessions/ci_abcdef/workspace/"
    # Use the session portion as a unique local subdir under /tmp
    clean = s3_prefix.strip("/").replace("/", "_")
    workdir = f"/tmp/{clean}"
    os.makedirs(workdir, exist_ok=True)
    return workdir


def _exec_python(event: Dict, s3_bucket: str, s3_prefix: str, timeout: int) -> Dict[str, Any]:
    code = event.get("code", "")
    script_name = event.get("script_name", "_lambda_exec.py")

    workdir = _get_session_workdir(s3_prefix)
    tmp_script = os.path.join("/tmp", script_name)

    # On a checkpoint resume, reconstitute the workdir from the durable S3
    # workspace (manifest + shards + prior outputs) BEFORE prefetch, so aa_env
    # continues from the last committed cursor. /tmp is never relied upon.
    if event.get("resume"):
        _sync_prefix_down(s3_bucket, s3_prefix, workdir)

    _prefetch_files(event.get("prefetch", []), s3_bucket, s3_prefix, workdir)

    print(f"[HANDLER] _exec_python: script={script_name}, code_len={len(code)}, workdir={workdir}")

    if code:
        with open(tmp_script, "w") as f:
            f.write(code)
        print(f"[HANDLER] Script written to {tmp_script}")
    elif not os.path.exists(tmp_script):
        return {"exit_code": 1, "stdout": "", "stderr": "No code provided and script not found"}

    # Make the aa_env checkpoint helper (shipped next to this handler) importable
    # by the user script, which runs with cwd=workdir. Prepend the handler's own
    # directory to PYTHONPATH so `import aa_env` resolves in any image layout.
    handler_dir = os.path.dirname(os.path.abspath(__file__))
    child_env = {**os.environ, "PYTHONUNBUFFERED": "1"}
    child_env["PYTHONPATH"] = handler_dir + (
        os.pathsep + child_env["PYTHONPATH"] if child_env.get("PYTHONPATH") else ""
    )

    result = subprocess.run(
        ["timeout", str(timeout), sys.executable, tmp_script],
        capture_output=True,
        cwd=workdir,
        env=child_env,
    )

    print(f"[HANDLER] _exec_python done: exit_code={result.returncode}, stdout_len={len(result.stdout)}, stderr_len={len(result.stderr)}")
    if result.returncode != 0:
        print(f"[HANDLER] stderr: {result.stderr[:500]}")

    # Upload script itself plus any output files to S3
    if s3_bucket and s3_prefix:
        # Upload the script
        try:
            s3_key = f"{s3_prefix.rstrip('/')}/{script_name}"
            with open(tmp_script, "rb") as f:
                _s3_client().put_object(Bucket=s3_bucket, Key=s3_key, Body=f.read())
            print(f"[HANDLER] Script uploaded to s3://{s3_bucket}/{s3_key}")
        except Exception as e:
            print(f"[HANDLER] WARNING: script upload failed: {e}")
        # Upload all output files from workdir
        _upload_workdir_files(s3_bucket, s3_prefix, workdir)

    return {
        "exit_code": result.returncode,
        "stdout": _cap_output(result.stdout),
        "stderr": _cap_output(result.stderr),
    }


def _exec_shell(event: Dict, s3_bucket: str, s3_prefix: str, timeout: int) -> Dict[str, Any]:
    command = event.get("command", "")
    if not command:
        return {"exit_code": 1, "stdout": "", "stderr": "No command provided"}

    workdir = _get_session_workdir(s3_prefix)

    _prefetch_files(event.get("prefetch", []), s3_bucket, s3_prefix, workdir)

    result = subprocess.run(
        ["timeout", str(timeout), "sh", "-c", command],
        capture_output=True,
        cwd=workdir,
        env=os.environ.copy(),
    )

    # Upload any files created by the shell command
    if s3_bucket and s3_prefix:
        _upload_workdir_files(s3_bucket, s3_prefix, workdir)

    return {
        "exit_code": result.returncode,
        "stdout": _cap_output(result.stdout),
        "stderr": _cap_output(result.stderr),
    }


def _write_file(event: Dict, s3_bucket: str, s3_prefix: str) -> Dict[str, Any]:
    rel_path = event.get("path", "")
    content = event.get("content", "")
    if not rel_path:
        return {"exit_code": 1, "stdout": "", "stderr": "No path provided"}

    s3_key = f"{s3_prefix.rstrip('/')}/{rel_path.lstrip('/')}"
    try:
        _s3_client().put_object(
            Bucket=s3_bucket,
            Key=s3_key,
            Body=content.encode("utf-8"),
        )
        return {"exit_code": 0, "stdout": f"Written {len(content)} bytes to {rel_path}", "stderr": ""}
    except Exception as e:
        return {"exit_code": 1, "stdout": "", "stderr": f"Write failed: {e}"}


def _read_file(event: Dict, s3_bucket: str, s3_prefix: str) -> Dict[str, Any]:
    rel_path = event.get("path", "")
    if not rel_path:
        return {"exit_code": 1, "stdout": "", "stderr": "No path provided"}

    s3_key = f"{s3_prefix.rstrip('/')}/{rel_path.lstrip('/')}"
    try:
        resp = _s3_client().get_object(Bucket=s3_bucket, Key=s3_key)
        content = resp["Body"].read().decode("utf-8", errors="replace")
        return {"exit_code": 0, "stdout": content, "stderr": ""}
    except Exception as e:
        if "NoSuchKey" in str(e):
            return {"exit_code": 1, "stdout": "", "stderr": f"File not found: {rel_path}"}
        return {"exit_code": 1, "stdout": "", "stderr": f"Read failed: {e}"}


def _list_files(event: Dict, s3_bucket: str, s3_prefix: str) -> Dict[str, Any]:
    try:
        s3 = _s3_client()
        paginator = s3.get_paginator("list_objects_v2")
        pages = paginator.paginate(Bucket=s3_bucket, Prefix=s3_prefix)

        lines = []
        for page in pages:
            for obj in page.get("Contents", []):
                key = obj["Key"]
                rel = key[len(s3_prefix):]
                if not rel or rel.endswith("/"):
                    continue
                size = obj.get("Size", 0)
                lines.append(f"f {size} {rel}")

        return {"exit_code": 0, "stdout": "\n".join(lines), "stderr": ""}
    except Exception as e:
        return {"exit_code": 1, "stdout": "", "stderr": f"List failed: {e}"}