"""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 traceback
from typing import Any, Dict, List

import boto3

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


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 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}, timeout={timeout}")
    print(f"[HANDLER] uid={os.getuid()}, gid={os.getgid()}")

    if action in ("exec_python", "exec_shell"):
        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 _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)

    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"}

    result = subprocess.run(
        ["timeout", str(timeout), sys.executable, tmp_script],
        capture_output=True,
        cwd=workdir,
        env={**os.environ, "PYTHONUNBUFFERED": "1"},
    )

    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)

    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}"}