"""OCLAWMA Skill: AWS

Official AWS skill for OCLAWMA providing comprehensive cloud resource
management capabilities including EC2, S3, Lambda, CloudWatch, IAM, RDS,
and Route53 operations.
"""

from __future__ import annotations

import json
import os
from datetime import datetime, timedelta
from typing import Any

import boto3
from botocore.exceptions import ClientError
from oclawma.skills import LazySkill, SkillMetadata

__all__ = ["AwsSkill"]


class AwsSkill(LazySkill):
    """AWS cloud resource management skill.

    This skill provides tools for managing AWS resources across multiple
    services including EC2, S3, Lambda, CloudWatch, IAM, RDS, and Route53.

    Features:
    - EC2 instance management (list, start, stop, describe)
    - S3 operations (list, upload, download, delete)
    - Lambda function management
    - CloudWatch metrics and logs
    - IAM user and role listing
    - RDS instance information
    - Route53 hosted zone management

    Example:
        >>> from oclawma.skills import SkillRegistry
        >>> registry = SkillRegistry()
        >>>
        >>> # List EC2 instances
        >>> result = await registry.execute_tool(
        ...     "aws", "ec2_list_instances",
        ...     region="us-east-1"
        ... )
        >>>
        >>> # List S3 objects
        >>> result = await registry.execute_tool(
        ...     "aws", "s3_list_objects",
        ...     bucket="my-bucket",
        ...     prefix="data/"
        ... )
    """

    def __init__(self, metadata: SkillMetadata) -> None:
        """Initialize the AWS skill.

        Args:
            metadata: Skill metadata from the entry point.
        """
        super().__init__(metadata)
        self._clients: dict[str, Any] = {}
        self._default_region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")

    def _get_client(self, service: str, region: str | None = None) -> Any:
        """Get boto3 client for a service, caching the result.

        Args:
            service: AWS service name (e.g., ec2, s3)
            region: AWS region (defaults to AWS_DEFAULT_REGION or us-east-1)

        Returns:
            Boto3 client instance
        """
        region = region or self._default_region
        cache_key = f"{service}:{region}"

        if cache_key not in self._clients:
            self._clients[cache_key] = boto3.client(service, region_name=region)

        return self._clients[cache_key]

    def _handle_error(self, error: ClientError) -> dict[str, Any]:
        """Handle AWS API errors.

        Args:
            error: Boto3 ClientError

        Returns:
            Standardized error result
        """
        error_code = error.response["Error"]["Code"]
        error_message = error.response["Error"]["Message"]

        return {
            "success": False,
            "output": "",
            "error": f"AWS Error ({error_code}): {error_message}",
            "exit_code": 1
        }

    def _serialize_datetime(self, obj: Any) -> Any:
        """Serialize datetime objects for JSON.

        Args:
            obj: Object to serialize

        Returns:
            Serialized object
        """
        if isinstance(obj, datetime):
            return obj.isoformat()
        return obj

    def _load(self) -> None:
        """Load the skill's tools."""
        self._tools = {
            # EC2 operations
            "ec2_list_instances": self._ec2_list_instances,
            "ec2_start_instance": self._ec2_start_instance,
            "ec2_stop_instance": self._ec2_stop_instance,
            "ec2_describe_instance": self._ec2_describe_instance,

            # S3 operations
            "s3_list_buckets": self._s3_list_buckets,
            "s3_list_objects": self._s3_list_objects,
            "s3_upload": self._s3_upload,
            "s3_download": self._s3_download,
            "s3_delete_object": self._s3_delete_object,

            # Lambda operations
            "lambda_list_functions": self._lambda_list_functions,
            "lambda_invoke": self._lambda_invoke,
            "lambda_get_logs": self._lambda_get_logs,

            # CloudWatch operations
            "cloudwatch_get_metrics": self._cloudwatch_get_metrics,

            # IAM operations
            "iam_list_users": self._iam_list_users,
            "iam_list_roles": self._iam_list_roles,

            # RDS operations
            "rds_describe_instances": self._rds_describe_instances,

            # Route53 operations
            "route53_list_hosted_zones": self._route53_list_hosted_zones,

            # STS operations
            "sts_get_caller_identity": self._sts_get_caller_identity,
        }
        self._loaded = True

    # -------------------------------------------------------------------------
    # EC2 Operations
    # -------------------------------------------------------------------------

    async def _ec2_list_instances(
        self,
        region: str | None = None,
        filters: dict[str, str] | None = None,
        max_results: int | None = None
    ) -> dict[str, Any]:
        """List EC2 instances."""
        try:
            ec2 = self._get_client("ec2", region)

            kwargs: dict[str, Any] = {}

            if filters:
                kwargs["Filters"] = [
                    {"Name": k, "Values": [v] if isinstance(v, str) else v}
                    for k, v in filters.items()
                ]

            if max_results:
                kwargs["MaxResults"] = max_results

            response = ec2.describe_instances(**kwargs)

            # Extract instance information
            instances = []
            for reservation in response.get("Reservations", []):
                for instance in reservation.get("Instances", []):
                    instances.append({
                        "InstanceId": instance.get("InstanceId"),
                        "InstanceType": instance.get("InstanceType"),
                        "State": instance.get("State", {}).get("Name"),
                        "PublicIpAddress": instance.get("PublicIpAddress"),
                        "PrivateIpAddress": instance.get("PrivateIpAddress"),
                        "LaunchTime": self._serialize_datetime(instance.get("LaunchTime")),
                        "Name": next(
                            (tag["Value"] for tag in instance.get("Tags", [])
                             if tag["Key"] == "Name"),
                            None
                        )
                    })

            return {
                "success": True,
                "output": json.dumps({"instances": instances}, indent=2),
                "data": {"instances": instances},
                "count": len(instances)
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _ec2_start_instance(
        self,
        instance_id: str,
        region: str | None = None
    ) -> dict[str, Any]:
        """Start an EC2 instance."""
        try:
            ec2 = self._get_client("ec2", region)
            response = ec2.start_instances(InstanceIds=[instance_id])

            return {
                "success": True,
                "output": f"Starting instance {instance_id}",
                "data": response.get("StartingInstances", [])
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _ec2_stop_instance(
        self,
        instance_id: str,
        region: str | None = None,
        force: bool = False
    ) -> dict[str, Any]:
        """Stop an EC2 instance."""
        try:
            ec2 = self._get_client("ec2", region)
            response = ec2.stop_instances(
                InstanceIds=[instance_id],
                Force=force
            )

            return {
                "success": True,
                "output": f"Stopping instance {instance_id}",
                "data": response.get("StoppingInstances", [])
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _ec2_describe_instance(
        self,
        instance_id: str,
        region: str | None = None
    ) -> dict[str, Any]:
        """Get detailed information about an EC2 instance."""
        try:
            ec2 = self._get_client("ec2", region)
            response = ec2.describe_instances(InstanceIds=[instance_id])

            instances = []
            for reservation in response.get("Reservations", []):
                instances.extend(reservation.get("Instances", []))

            # Serialize datetime objects
            for instance in instances:
                for key, value in instance.items():
                    instance[key] = self._serialize_datetime(value)

            return {
                "success": True,
                "output": json.dumps({"instances": instances}, indent=2, default=str),
                "data": {"instances": instances}
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # S3 Operations
    # -------------------------------------------------------------------------

    async def _s3_list_buckets(self, max_results: int | None = None) -> dict[str, Any]:
        """List S3 buckets."""
        try:
            s3 = self._get_client("s3")
            response = s3.list_buckets()

            buckets = response.get("Buckets", [])
            if max_results:
                buckets = buckets[:max_results]

            # Format bucket info
            bucket_list = []
            for bucket in buckets:
                bucket_list.append({
                    "Name": bucket.get("Name"),
                    "CreationDate": self._serialize_datetime(bucket.get("CreationDate"))
                })

            return {
                "success": True,
                "output": json.dumps({"buckets": bucket_list}, indent=2),
                "data": {"buckets": bucket_list},
                "count": len(bucket_list)
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _s3_list_objects(
        self,
        bucket: str,
        prefix: str | None = None,
        max_results: int | None = None
    ) -> dict[str, Any]:
        """List objects in an S3 bucket."""
        try:
            s3 = self._get_client("s3")

            kwargs: dict[str, Any] = {"Bucket": bucket}
            if prefix:
                kwargs["Prefix"] = prefix
            if max_results:
                kwargs["MaxKeys"] = max_results

            response = s3.list_objects_v2(**kwargs)

            objects = []
            for obj in response.get("Contents", []):
                objects.append({
                    "Key": obj.get("Key"),
                    "Size": obj.get("Size"),
                    "LastModified": self._serialize_datetime(obj.get("LastModified")),
                    "StorageClass": obj.get("StorageClass")
                })

            return {
                "success": True,
                "output": json.dumps({"objects": objects}, indent=2),
                "data": {"objects": objects, "prefix": prefix},
                "count": len(objects)
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _s3_upload(
        self,
        bucket: str,
        key: str,
        file_path: str,
        metadata: dict[str, str] | None = None
    ) -> dict[str, Any]:
        """Upload a file to S3."""
        try:
            s3 = self._get_client("s3")

            kwargs: dict[str, Any] = {
                "Bucket": bucket,
                "Key": key,
                "Filename": file_path
            }

            if metadata:
                kwargs["ExtraArgs"] = {"Metadata": metadata}

            s3.upload_file(**kwargs)

            return {
                "success": True,
                "output": f"Uploaded {file_path} to s3://{bucket}/{key}",
                "data": {"bucket": bucket, "key": key}
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _s3_download(
        self,
        bucket: str,
        key: str,
        file_path: str
    ) -> dict[str, Any]:
        """Download a file from S3."""
        try:
            s3 = self._get_client("s3")
            s3.download_file(bucket, key, file_path)

            return {
                "success": True,
                "output": f"Downloaded s3://{bucket}/{key} to {file_path}",
                "data": {"bucket": bucket, "key": key, "file_path": file_path}
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _s3_delete_object(
        self,
        bucket: str,
        key: str
    ) -> dict[str, Any]:
        """Delete an object from S3."""
        try:
            s3 = self._get_client("s3")
            s3.delete_object(Bucket=bucket, Key=key)

            return {
                "success": True,
                "output": f"Deleted s3://{bucket}/{key}",
                "data": {"bucket": bucket, "key": key}
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # Lambda Operations
    # -------------------------------------------------------------------------

    async def _lambda_list_functions(
        self,
        region: str | None = None,
        max_results: int | None = None
    ) -> dict[str, Any]:
        """List Lambda functions."""
        try:
            lambda_client = self._get_client("lambda", region)

            kwargs: dict[str, Any] = {}
            if max_results:
                kwargs["MaxItems"] = max_results

            response = lambda_client.list_functions(**kwargs)

            functions = []
            for func in response.get("Functions", []):
                functions.append({
                    "FunctionName": func.get("FunctionName"),
                    "Runtime": func.get("Runtime"),
                    "Handler": func.get("Handler"),
                    "State": func.get("State"),
                    "LastModified": func.get("LastModified"),
                    "MemorySize": func.get("MemorySize"),
                    "Timeout": func.get("Timeout")
                })

            return {
                "success": True,
                "output": json.dumps({"functions": functions}, indent=2),
                "data": {"functions": functions},
                "count": len(functions)
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _lambda_invoke(
        self,
        function_name: str,
        payload: dict[str, Any] | None = None,
        region: str | None = None
    ) -> dict[str, Any]:
        """Invoke a Lambda function."""
        try:
            lambda_client = self._get_client("lambda", region)

            kwargs: dict[str, Any] = {
                "FunctionName": function_name,
                "InvocationType": "RequestResponse"
            }

            if payload:
                kwargs["Payload"] = json.dumps(payload)

            response = lambda_client.invoke(**kwargs)

            # Read and parse response payload
            response_payload = response["Payload"].read().decode("utf-8")
            try:
                response_data = json.loads(response_payload)
            except json.JSONDecodeError:
                response_data = response_payload

            return {
                "success": response["StatusCode"] == 200,
                "output": json.dumps({"response": response_data}, indent=2),
                "data": {
                    "status_code": response["StatusCode"],
                    "response": response_data
                }
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _lambda_get_logs(
        self,
        function_name: str,
        region: str | None = None,
        start_time: str | None = None,
        end_time: str | None = None,
        limit: int = 100
    ) -> dict[str, Any]:
        """Get CloudWatch logs for a Lambda function."""
        try:
            logs = self._get_client("logs", region)

            # Construct log group name from function name
            log_group = f"/aws/lambda/{function_name}"

            kwargs: dict[str, Any] = {
                "logGroupName": log_group,
                "limit": limit,
                "startFromHead": False
            }

            if start_time:
                kwargs["startTime"] = int(datetime.fromisoformat(start_time).timestamp() * 1000)
            if end_time:
                kwargs["endTime"] = int(datetime.fromisoformat(end_time).timestamp() * 1000)

            response = logs.filter_log_events(**kwargs)

            events = []
            for event in response.get("events", []):
                events.append({
                    "timestamp": datetime.fromtimestamp(
                        event["timestamp"] / 1000
                    ).isoformat(),
                    "message": event["message"],
                    "logStreamName": event["logStreamName"]
                })

            return {
                "success": True,
                "output": json.dumps({"events": events}, indent=2),
                "data": {"events": events},
                "count": len(events)
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # CloudWatch Operations
    # -------------------------------------------------------------------------

    async def _cloudwatch_get_metrics(
        self,
        namespace: str,
        metric_name: str,
        dimensions: dict[str, str] | None = None,
        start_time: str | None = None,
        end_time: str | None = None,
        period: int = 300,
        region: str | None = None
    ) -> dict[str, Any]:
        """Get CloudWatch metrics."""
        try:
            cloudwatch = self._get_client("cloudwatch", region)

            # Default to last hour if not specified
            end = datetime.utcnow()
            start = end - timedelta(hours=1)

            if end_time:
                end = datetime.fromisoformat(end_time)
            if start_time:
                start = datetime.fromisoformat(start_time)

            kwargs: dict[str, Any] = {
                "Namespace": namespace,
                "MetricName": metric_name,
                "StartTime": start,
                "EndTime": end,
                "Period": period,
                "Statistics": ["Average", "Minimum", "Maximum", "Sum", "SampleCount"]
            }

            if dimensions:
                kwargs["Dimensions"] = [
                    {"Name": k, "Value": v}
                    for k, v in dimensions.items()
                ]

            response = cloudwatch.get_metric_statistics(**kwargs)

            datapoints = []
            for dp in response.get("Datapoints", []):
                datapoints.append({
                    "Timestamp": self._serialize_datetime(dp.get("Timestamp")),
                    "Average": dp.get("Average"),
                    "Minimum": dp.get("Minimum"),
                    "Maximum": dp.get("Maximum"),
                    "Sum": dp.get("Sum"),
                    "SampleCount": dp.get("SampleCount"),
                    "Unit": dp.get("Unit")
                })

            # Sort by timestamp
            datapoints.sort(key=lambda x: x["Timestamp"])

            return {
                "success": True,
                "output": json.dumps({"datapoints": datapoints}, indent=2),
                "data": {"datapoints": datapoints, "label": response.get("Label")},
                "count": len(datapoints)
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # IAM Operations
    # -------------------------------------------------------------------------

    async def _iam_list_users(self, max_results: int | None = None) -> dict[str, Any]:
        """List IAM users."""
        try:
            iam = self._get_client("iam")

            kwargs: dict[str, Any] = {}
            if max_results:
                kwargs["MaxItems"] = max_results

            response = iam.list_users(**kwargs)

            users = []
            for user in response.get("Users", []):
                users.append({
                    "UserName": user.get("UserName"),
                    "UserId": user.get("UserId"),
                    "Arn": user.get("Arn"),
                    "CreateDate": self._serialize_datetime(user.get("CreateDate"))
                })

            return {
                "success": True,
                "output": json.dumps({"users": users}, indent=2),
                "data": {"users": users},
                "count": len(users)
            }

        except ClientError as e:
            return self._handle_error(e)

    async def _iam_list_roles(self, max_results: int | None = None) -> dict[str, Any]:
        """List IAM roles."""
        try:
            iam = self._get_client("iam")

            kwargs: dict[str, Any] = {}
            if max_results:
                kwargs["MaxItems"] = max_results

            response = iam.list_roles(**kwargs)

            roles = []
            for role in response.get("Roles", []):
                roles.append({
                    "RoleName": role.get("RoleName"),
                    "RoleId": role.get("RoleId"),
                    "Arn": role.get("Arn"),
                    "CreateDate": self._serialize_datetime(role.get("CreateDate"))
                })

            return {
                "success": True,
                "output": json.dumps({"roles": roles}, indent=2),
                "data": {"roles": roles},
                "count": len(roles)
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # RDS Operations
    # -------------------------------------------------------------------------

    async def _rds_describe_instances(
        self,
        region: str | None = None,
        max_results: int | None = None
    ) -> dict[str, Any]:
        """List RDS instances."""
        try:
            rds = self._get_client("rds", region)

            kwargs: dict[str, Any] = {}
            if max_results:
                kwargs["MaxRecords"] = max_results

            response = rds.describe_db_instances(**kwargs)

            instances = []
            for instance in response.get("DBInstances", []):
                instances.append({
                    "DBInstanceIdentifier": instance.get("DBInstanceIdentifier"),
                    "DBInstanceClass": instance.get("DBInstanceClass"),
                    "Engine": instance.get("Engine"),
                    "DBInstanceStatus": instance.get("DBInstanceStatus"),
                    "Endpoint": instance.get("Endpoint", {}).get("Address"),
                    "AllocatedStorage": instance.get("AllocatedStorage")
                })

            return {
                "success": True,
                "output": json.dumps({"instances": instances}, indent=2),
                "data": {"instances": instances},
                "count": len(instances)
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # Route53 Operations
    # -------------------------------------------------------------------------

    async def _route53_list_hosted_zones(
        self,
        max_results: int | None = None
    ) -> dict[str, Any]:
        """List Route53 hosted zones."""
        try:
            route53 = self._get_client("route53")

            kwargs: dict[str, Any] = {}
            if max_results:
                kwargs["MaxItems"] = str(max_results)

            response = route53.list_hosted_zones(**kwargs)

            zones = []
            for zone in response.get("HostedZones", []):
                zones.append({
                    "Id": zone.get("Id"),
                    "Name": zone.get("Name"),
                    "ResourceRecordSetCount": zone.get("ResourceRecordSetCount"),
                    "PrivateZone": zone.get("Config", {}).get("PrivateZone")
                })

            return {
                "success": True,
                "output": json.dumps({"hosted_zones": zones}, indent=2),
                "data": {"hosted_zones": zones},
                "count": len(zones)
            }

        except ClientError as e:
            return self._handle_error(e)

    # -------------------------------------------------------------------------
    # STS Operations
    # -------------------------------------------------------------------------

    async def _sts_get_caller_identity(self) -> dict[str, Any]:
        """Get current AWS identity."""
        try:
            sts = self._get_client("sts")
            response = sts.get_caller_identity()

            identity = {
                "Account": response.get("Account"),
                "Arn": response.get("Arn"),
                "UserId": response.get("UserId")
            }

            return {
                "success": True,
                "output": json.dumps(identity, indent=2),
                "data": identity
            }

        except ClientError as e:
            return self._handle_error(e)
