"""Lambda service emulator."""
import base64, json, logging, os, subprocess, sys, tempfile, time, uuid, zipfile
from dataclasses import dataclass, field
from typing import Optional
from flask import Request, Response
from localrun.config import get_config
from localrun.utils import error_response, iso_timestamp, new_request_id

logger = logging.getLogger("localrun.lambda")

def _json_resp(data, status=200):
    return Response(json.dumps(data, default=str), status=status, content_type="application/json")

@dataclass
class LambdaFunction:
    name: str; arn: str; runtime: str; handler: str; role: str
    code_zip: bytes = b""; description: str = ""; timeout: int = 30
    memory_size: int = 128; environment: dict = field(default_factory=dict)
    last_modified: str = ""; code_sha256: str = ""; code_size: int = 0
    def __post_init__(self):
        if not self.last_modified: self.last_modified = iso_timestamp()
        import hashlib
        self.code_sha256 = base64.b64encode(hashlib.sha256(self.code_zip).digest()).decode()
        self.code_size = len(self.code_zip)

class LambdaService:
    def __init__(self):
        self.functions: dict[str, LambdaFunction] = {}

    def handle(self, req: Request, path: str) -> Response:
        method = req.method
        # Lambda REST API paths: 2015-03-31/functions[/name[/invocations|/code|/configuration]]
        parts = path.strip("/").split("/")
        # Strip version prefix
        if parts and parts[0].startswith("2015") or parts[0].startswith("2014"):
            parts = parts[1:]
        if not parts: return _json_resp({"error": "Invalid path"}, 400)
        resource = parts[0]  # "functions"
        if resource == "functions":
            if len(parts) == 1:
                if method == "GET": return self._list_functions(req)
                if method == "POST": return self._create_function(req)
            elif len(parts) >= 2:
                func_name = parts[1]
                if len(parts) == 2:
                    if method == "GET": return self._get_function(func_name)
                    if method == "DELETE": return self._delete_function(func_name)
                elif len(parts) >= 3:
                    sub = parts[2]
                    if sub == "invocations" and method == "POST":
                        return self._invoke(req, func_name)
                    if sub == "code" and method == "PUT":
                        return self._update_code(req, func_name)
                    if sub == "configuration":
                        if method == "GET": return self._get_function(func_name)
                        if method == "PUT": return self._update_config(req, func_name)
        return _json_resp({"message": "Not Found"}, 404)

    def _func_config(self, f: LambdaFunction) -> dict:
        return {
            "FunctionName": f.name, "FunctionArn": f.arn, "Runtime": f.runtime,
            "Handler": f.handler, "Role": f.role, "Description": f.description,
            "Timeout": f.timeout, "MemorySize": f.memory_size,
            "CodeSha256": f.code_sha256, "CodeSize": f.code_size,
            "LastModified": f.last_modified, "State": "Active",
            "Environment": {"Variables": f.environment},
        }

    def _create_function(self, req):
        try:
            body = json.loads(req.get_data(as_text=True) or "{}")
        except Exception:
            body = {}
        name = body.get("FunctionName", "")
        if not name: return _json_resp({"message": "FunctionName required"}, 400)
        if name in self.functions: return _json_resp({"message": f"Function {name} already exists"}, 409)
        c = get_config()
        code = body.get("Code", {}); zip_bytes = b""
        if "ZipFile" in code: zip_bytes = base64.b64decode(code["ZipFile"])
        fn = LambdaFunction(
            name=name, arn=f"arn:aws:lambda:{c.region}:{c.account_id}:function:{name}",
            runtime=body.get("Runtime", "python3.12"), handler=body.get("Handler", "index.handler"),
            role=body.get("Role", ""), code_zip=zip_bytes, description=body.get("Description", ""),
            timeout=body.get("Timeout", 30), memory_size=body.get("MemorySize", 128),
            environment=body.get("Environment", {}).get("Variables", {}),
        )
        self.functions[name] = fn
        logger.info("Created function: %s", name)
        return _json_resp(self._func_config(fn), 201)

    def _list_functions(self, req):
        funcs = [self._func_config(f) for f in self.functions.values()]
        return _json_resp({"Functions": funcs})

    def _get_function(self, name):
        fn = self.functions.get(name)
        if not fn: return _json_resp({"message": f"Function {name} not found"}, 404)
        return _json_resp({"Configuration": self._func_config(fn), "Code": {"RepositoryType": "S3"}})

    def _delete_function(self, name):
        if name not in self.functions: return _json_resp({"message": f"Function {name} not found"}, 404)
        del self.functions[name]; logger.info("Deleted function: %s", name)
        return Response("", 204)

    def _update_code(self, req, name):
        fn = self.functions.get(name)
        if not fn: return _json_resp({"message": f"Function {name} not found"}, 404)
        try:
            body = json.loads(req.get_data(as_text=True) or "{}")
        except Exception:
            body = {}
        if "ZipFile" in body: fn.code_zip = base64.b64decode(body["ZipFile"])
        fn.last_modified = iso_timestamp()
        import hashlib
        fn.code_sha256 = base64.b64encode(hashlib.sha256(fn.code_zip).digest()).decode()
        fn.code_size = len(fn.code_zip)
        return _json_resp(self._func_config(fn))

    def _update_config(self, req, name):
        fn = self.functions.get(name)
        if not fn: return _json_resp({"message": f"Function {name} not found"}, 404)
        try:
            body = json.loads(req.get_data(as_text=True) or "{}")
        except Exception:
            body = {}
        if "Handler" in body: fn.handler = body["Handler"]
        if "Runtime" in body: fn.runtime = body["Runtime"]
        if "Timeout" in body: fn.timeout = body["Timeout"]
        if "MemorySize" in body: fn.memory_size = body["MemorySize"]
        if "Description" in body: fn.description = body["Description"]
        if "Environment" in body: fn.environment = body["Environment"].get("Variables", {})
        fn.last_modified = iso_timestamp()
        return _json_resp(self._func_config(fn))

    def _invoke(self, req, name):
        fn = self.functions.get(name)
        if not fn: return _json_resp({"message": f"Function {name} not found"}, 404)
        payload = req.get_data(as_text=True) or "{}"
        invocation_type = req.headers.get("X-Amz-Invocation-Type", "RequestResponse")
        # Execute the function
        result = self._execute_function(fn, payload)
        if invocation_type == "Event":
            return Response("", 202)
        resp = Response(json.dumps(result.get("body", {}), default=str), 200, content_type="application/json")
        if result.get("error"): resp.headers["X-Amz-Function-Error"] = "Unhandled"
        return resp

    def _execute_function(self, fn: LambdaFunction, payload: str) -> dict:
        """Execute a Lambda function by extracting zip and running the handler."""
        if not fn.code_zip:
            return {"error": True, "body": {"errorMessage": "No function code"}}
        try:
            with tempfile.TemporaryDirectory() as tmpdir:
                # Extract zip
                zip_path = os.path.join(tmpdir, "code.zip")
                with open(zip_path, "wb") as f: f.write(fn.code_zip)
                with zipfile.ZipFile(zip_path, "r") as zf: zf.extractall(tmpdir)
                # Parse handler: module.function
                module_name, func_name = fn.handler.rsplit(".", 1)
                # Build runner script
                runner = f"""
import sys, json, os
sys.path.insert(0, '{tmpdir}')
os.environ.update({json.dumps(fn.environment)})
import {module_name}
event = json.loads('''{payload}''')
context = type('Context', (), {{'function_name': '{fn.name}', 'memory_limit_in_mb': {fn.memory_size}, 'invoked_function_arn': '{fn.arn}', 'aws_request_id': '{str(uuid.uuid4())}'}})()
result = {module_name}.{func_name}(event, context)
print(json.dumps(result, default=str))
"""
                result = subprocess.run(
                    [sys.executable, "-c", runner], capture_output=True, text=True,
                    timeout=fn.timeout, cwd=tmpdir
                )
                if result.returncode != 0:
                    return {"error": True, "body": {"errorMessage": result.stderr.strip()}}
                try:
                    return {"body": json.loads(result.stdout.strip())}
                except json.JSONDecodeError:
                    return {"body": result.stdout.strip()}
        except subprocess.TimeoutExpired:
            return {"error": True, "body": {"errorMessage": f"Task timed out after {fn.timeout} seconds"}}
        except Exception as e:
            return {"error": True, "body": {"errorMessage": str(e)}}
