import os
import json
import argparse
import sys
from config import SPLITIO_WORKSPACE_ID, SPLITIO_ADMIN_API_KEY
from splitio_client import SplitIOClient

import boto3

def upload_to_s3(bucket, key, data):
    s3 = boto3.client("s3")
    s3.put_object(Bucket=bucket, Key=key, Body=json.dumps(data, indent=2))

    
def download_from_s3(bucket, key):
    s3 = boto3.client("s3")
    response = s3.get_object(Bucket=bucket, Key=key)
    return json.loads(response["Body"].read())


def backup_feature_flag(splitio, flag_name, backup_dir="backups", s3_bucket=None, only_s3=False, dry_run=False):
    flag_dir = os.path.join(backup_dir, flag_name)
    if not only_s3 and not dry_run:
        os.makedirs(os.path.join(flag_dir, "definitions"), exist_ok=True)

    metadata = splitio.get_feature_flag_metadata(flag_name)
    environments = splitio.get_environments()

    if dry_run:
        print(f"\n[DRY-RUN] Would back up feature flag '{flag_name}'")
        if s3_bucket:
            print(f"  - S3: s3://{s3_bucket}/{backup_dir}/{flag_name}/metadata.json")
        if not only_s3:
            print(f"  - Local: {os.path.join(flag_dir, 'metadata.json')}")
    else:
        if s3_bucket:
            upload_to_s3(s3_bucket, f"{backup_dir}/{flag_name}/metadata.json", metadata)
        if not only_s3:
            with open(os.path.join(flag_dir, "metadata.json"), "w") as f:
                json.dump(metadata, f, indent=2)

    for env in environments:
        env_id = env["id"]
        env_name = env["name"]
        definition = splitio.get_feature_flag_definition(env_id, flag_name)
        if not definition:
            continue

        if dry_run:
            if s3_bucket:
                print(f"  - S3: s3://{s3_bucket}/{backup_dir}/{flag_name}/definitions/{env_name}.json")
            if not only_s3:
                print(f"  - Local: {os.path.join(flag_dir, 'definitions', f'{env_name}.json')}")
        else:
            if s3_bucket:
                upload_to_s3(s3_bucket, f"{backup_dir}/{flag_name}/definitions/{env_name}.json", definition)
            if not only_s3:
                path = os.path.join(flag_dir, "definitions", f"{env_name}.json")
                with open(path, "w") as f:
                    json.dump(definition, f, indent=2)

    if not dry_run:
        print(f"📦 Backup complete for '{flag_name}'")
    else:
        print(f"[DRY-RUN] Backup simulation completed for '{flag_name}'")

def try_backup_flags(splitio, ff_names, s3_bucket=None, only_s3=False, dry_run=False):
    backed_up, failed = [], []

    for name in ff_names:
        name = name.strip()
        if not name:
            continue
        try:
            backup_feature_flag(
                splitio,
                name,
                s3_bucket=s3_bucket,
                only_s3=only_s3,
                dry_run=dry_run
            )
            backed_up.append(name)
        except Exception as e:
            prefix = "[DRY-RUN] ⚠️" if dry_run else "⚠️"
            print(f"{prefix} Backup failed for '{name}': {e}")
            failed.append((name, e))

    return backed_up, failed


def bulk_delete_feature_flags(splitio, ff_names, s3_bucket=None, only_s3=False, dry_run=False, skip_backup=False):
    success, failed = [], []
    
    if skip_backup:
        backed_up = [name.strip() for name in ff_names if name.strip()]
        failed_backup = []
    else:
        backed_up, failed_backup = try_backup_flags(splitio, ff_names, s3_bucket, only_s3, dry_run)

    for name in backed_up:        
        if dry_run:
            print(f"[DRY-RUN] Would delete feature flag: '{name}'")
            success.append(name)
            continue

        try:
            splitio.delete_feature_flag(name)
            print(f"🗑️ Deleted: {name}")
            success.append(name)
        except Exception as e:
            print(f"❌ Error deleting {name}: {e}")
            failed.append((name, str(e)))
    
    print("\n✅ Summary:")
    if dry_run:
        print(f"  [DRY-RUN] Total processed: {len(ff_names)}")
        print(f"  [DRY-RUN] Would delete: {len(success)}")
        print(f"  [DRY-RUN] Backup failed/skipped: {len(failed) + len(failed_backup)}")
    else:
        print(f"  Total processed: {len(ff_names)}")
        print(f"  Deleted: {len(success)}")
        print(f"  Failed: {len(failed) + len(failed_backup)}")
    
    if failed or failed_backup:
        print("⚠️ Errors:")
        for name, reason in failed + failed_backup:
            print(f"    - {name}: {reason}")


def backup_multiple_flags(splitio, ff_names, s3_bucket=None, only_s3=False, dry_run=False):
    backed_up, failed = try_backup_flags(splitio, ff_names, s3_bucket, only_s3, dry_run)
    message = f"\n[DRY-RUN]✅ Backup completed. Success: {len(backed_up)}, Failed: {len(failed)}" if dry_run else f"\n✅ Backup completed. Success: {len(backed_up)}, Failed: {len(failed)}"
    print(message)

def restore_feature_flag(splitio, flag_name, backup_dir="backups", s3_bucket=None, dry_run=False):
    # Load metadata
    if s3_bucket:
        s3_key = f"{backup_dir}/{flag_name}/metadata.json"
        try:
            metadata = download_from_s3(s3_bucket, s3_key)
            print(f"{'[DRY-RUN] ' if dry_run else ''}✅ Loaded metadata from: s3://{s3_bucket}/{s3_key}")
        except Exception:
            print(f"❌ Backup metadata not found in S3 at 's3://{s3_bucket}/{s3_key}' for '{flag_name}'")
            return          
    else:
        metadata_path = os.path.join(backup_dir, flag_name, "metadata.json")
        if not os.path.exists(metadata_path):
            print(f"❌ Backup metadata not found locally at '{metadata_path}' for '{flag_name}'")
            return
        with open(metadata_path, "r") as f:
            metadata = json.load(f)
        print(f"{'[DRY-RUN] ' if dry_run else ''}✅ Loaded metadata from: {metadata_path}")

    # Fetch traffic type ID
    traffic_type_name = metadata.get("trafficType", {}).get("name", "unknown")
    traffic_types = splitio.get_traffic_types()
    traffic_type_id = next((t["id"] for t in traffic_types if t["name"] == traffic_type_name), None)
    if not traffic_type_id:
        print(f"❌ Traffic type '{traffic_type_name}' not found.")
        return

    if dry_run:
        print(f"[DRY-RUN] ✅ Would recreate feature flag '{flag_name}' with traffic type '{traffic_type_name}'")
    else:
        try:
            splitio.recreate_feature_flag(flag_name, traffic_type_id, metadata)
            print(f"✅ Restored feature flag '{flag_name}'.")
        except Exception as e:
            if "409" in str(e):
                print(f"⚠️ Feature flag '{flag_name}' already exists.")
            else:
                print(f"❌ Failed to recreate FF '{flag_name}': {e}")
                return

    # Prepare patch operations
    patch_ops = []
    for i, tag in enumerate(metadata.get("tags", [])):
        patch_ops.append({"op": "add", "path": f"/tags/{i}", "value": {"name": tag["name"]}})

    if metadata.get("rolloutStatus", {}).get("id"):
        patch_ops.append({"op": "replace", "path": "/rolloutStatus/id", "value": metadata["rolloutStatus"]["id"]})

    if metadata.get("description"):
        patch_ops.append({"op": "replace", "path": "/description", "value": metadata["description"]})

        
    if dry_run:
        print(f"[DRY-RUN] ✅ Would patch feature flag '{flag_name}' with tags/rollout status/description")
    else:
        try:
            splitio.patch_feature_flag(flag_name, patch_ops)
            print(f"🩹 Patched tags/rollout status for '{flag_name}'.")
        except Exception as e:
            print(f"⚠️ Failed to patch flag '{flag_name}': {e}")
            return


    environments = splitio.get_environments()
    for env in environments:
        env_name = env["name"]
        env_id = env["id"]

        if s3_bucket:
            s3_key = f"{backup_dir}/{flag_name}/definitions/{env_name}.json"
            try:
                definition = download_from_s3(s3_bucket, s3_key)
                print(f"{'[DRY-RUN] ' if dry_run else ''}  ✅ Loaded definition from: s3://{s3_bucket}/{s3_key}")
            except Exception:
                print(f"  ⚠️ No definition found in S3 for environment '{env_name}'")
                continue
        else:
            env_def_path = os.path.join(backup_dir, flag_name, "definitions", f"{env_name}.json")
            if not os.path.exists(env_def_path):
                print(f"  ⚠️ No definition found for environment '{env_name}'")
                continue
            with open(env_def_path, "r") as f:
                definition = json.load(f)
            print(f"{'[DRY-RUN] ' if dry_run else ''}  ✅ Loaded definition from: {env_def_path}")

        definition_data = {
            "treatments": definition.get("treatments", []),
            "defaultTreatment": definition.get("defaultTreatment", "off"),
            "baselineTreatment": definition.get("baselineTreatment", "off"),
            "rules": definition.get("rules", []),
            "trafficAllocation": definition.get("trafficAllocation", 100),
            "defaultRule": definition.get("defaultRule", []),
        }

        if dry_run:
            print(f"[DRY-RUN]   ✅ Would restore definition for environment '{env_name}'")
        else:
            try:
                splitio.post_definition_to_env(flag_name, env_id, definition_data)
                print(f"  🔁 Restored definition for '{env_name}'")
            except Exception as e:
                print(f"  ❌ Failed to restore definition for '{env_name}': {e}")


def process_feature_flags_from_args(args):
    if args.file and not os.path.exists(args.file):
        print(f"❌ File '{args.file}' does not exist.")
        sys.exit(1)
    
    if args.file:
        with open(args.file, "r") as f:
            ff_names = f.readlines()
    else:
        ff_names = args.names.split(",")
    return [name.strip() for name in ff_names if name.strip()]


def load_args():
    parser = argparse.ArgumentParser(description="Split.io Feature Flag CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # --- bulk-delete ---
    bulk_parser = subparsers.add_parser("bulk-delete", help="Bulk delete feature flags")
    group = bulk_parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--file", help="Text file with feature flag names (one per line)")
    group.add_argument("--names", help="Comma-separated feature flag names")
    bulk_parser.add_argument("--s3-backup-bucket", help="S3 bucket name to store backups")
    bulk_parser.add_argument("--only-s3", action="store_true", help="Do not write backups to local disk")
    bulk_parser.add_argument("--dry-run", action="store_true", help="Print actions without deleting flags")
    bulk_parser.add_argument("--skip-backup", action="store_true", help="Do not create backup before deleting")

    # --- restore ---
    restore_parser = subparsers.add_parser("restore", help="Restore feature flag from backup")
    restore_parser.add_argument("--name", required=True, help="Feature flag name")
    restore_parser.add_argument("--s3-backup-bucket", help="S3 bucket to restore from")
    restore_parser.add_argument("--dry-run", action="store_true", help="Print what would happen without making changes")
    
     # --- backup ---
    backup_parser = subparsers.add_parser("backup", help="Backup feature flags only")
    group = backup_parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--file", help="Text file with feature flag names (one per line)")
    group.add_argument("--names", help="Comma-separated feature flag names")
    backup_parser.add_argument("--s3-backup-bucket", help="S3 bucket name to store backups")
    backup_parser.add_argument("--only-s3", action="store_true", help="Do not write backups to local disk")
    backup_parser.add_argument("--dry-run", action="store_true", help="Print what would be backed up")

    return parser.parse_args()


def main() -> None:
    """Entrypoint."""
    args = load_args()
    splitio = SplitIOClient()
    
    if args.command == "bulk-delete":
        ff_names = process_feature_flags_from_args(args)
        bulk_delete_feature_flags(
            splitio,
            ff_names,
            s3_bucket=args.s3_backup_bucket,
            only_s3=args.only_s3,
            dry_run=args.dry_run,
            skip_backup=args.skip_backup,
        )

    elif args.command == "restore":
        restore_feature_flag(
            splitio,
            args.name,
            s3_bucket=args.s3_backup_bucket,
            dry_run=args.dry_run,
        )
    
    elif args.command == "backup":
        ff_names = process_feature_flags_from_args(args)
        backup_multiple_flags(
            splitio,
            ff_names,
            s3_bucket=args.s3_backup_bucket,
            only_s3=args.only_s3,
            dry_run=args.dry_run,
        )


if __name__ == '__main__':
    if not SPLITIO_ADMIN_API_KEY or not SPLITIO_WORKSPACE_ID:
        print("❌ Please set SPLIT_ADMIN_API_KEY and WORKSPACE_ID in .env or environment variables.")
        sys.exit(1)

    main()