import time

import boto3
import csv
import os
import requests
import json
from collections import defaultdict

from datadog_api_client import ApiClient, Configuration
from datadog_api_client.v2.api.software_catalog_api import SoftwareCatalogApi

OUTPUT_CSV = "ecs_service_catalog.csv"
GITHUB_TOKEN = os.getenv("GITHUB_TOKEN", "dummy")
GITHUB_ORG = "theorchard"
JENKINS_ORG_URL = "https://pipeline.theorchard.io/job/theorchard/api/json"
JENKINS_USER = os.getenv("JENKINS_USER", "dummy")
JENKINS_API_TOKEN = os.getenv("JENKINS_API_TOKEN", "dummy")
ENVIRONMENTS = ["dev", "qa", "stage", "uat", "prod", "shared", "backup", "networking", "unknown"]

iam_client = boto3.client('iam')
sts_client = boto3.client('sts')
ssm_client = boto3.client('ssm')
dd_config = Configuration()

REPO_OVERRIDES_PATH = os.path.join(os.path.dirname(__file__), 'ecs_repo_overrides.json')
if os.path.exists(REPO_OVERRIDES_PATH):
    with open(REPO_OVERRIDES_PATH, 'r') as f:
        REPO_OVERRIDES = json.load(f)
else:
    REPO_OVERRIDES = {}

def get_client(username, account, service):
    if account['name'] != 'prod':
        assume_role_response = sts_client.assume_role(
            RoleArn=f"arn:aws:iam::{account['account_id']}:role/{account['role']}",
            RoleSessionName=username
        )
        credentials = assume_role_response["Credentials"]
        print(f"Assumed role for account {account['account_id']} ({account['name']}) successfully.")
        return boto3.client(
            service,
            aws_access_key_id=credentials["AccessKeyId"],
            aws_secret_access_key=credentials["SecretAccessKey"],
            aws_session_token=credentials["SessionToken"],
        )
    else:
        return boto3.client(service)

def get_aws_accounts():
    aws_accounts = []

    for page in ssm_client.get_paginator('describe_parameters').paginate(
            Shared=True,
            ParameterFilters=[
                {
                    'Key': 'Name',
                    'Option': 'BeginsWith',
                    'Values': [
                        '/shared/aws-account-ids/',
                    ]
                },
            ]
    ):
            get_parameters_response = ssm_client.get_parameters(
            Names=[param['ARN'] for param in page['Parameters']],
            )
            aws_accounts.extend([json.loads(parameter['Value']) for parameter in get_parameters_response['Parameters']])

    return aws_accounts

def get_repo_from_commit(commit_sha):
    url = f"https://api.github.com/search/commits?q={commit_sha}+org:{GITHUB_ORG}"
    headers = {
        "Accept": "application/vnd.github.cloak-preview",
        "Authorization": f"token {GITHUB_TOKEN}"
    }
    max_retries = 10
    retry_counter = 0
    do_retry = True
    while do_retry and retry_counter < max_retries:
        resp = requests.get(url, headers=headers)
        # Throttle to avoid hitting the 30 req/min rate limit
        time.sleep(2)
        if resp.status_code == 200:
            items = resp.json().get("items", [])
            if items:
                return items[0]["repository"]["full_name"]
            else:
                # Search API is best effort so will occasionally return no items for a valid commit, so try again.
                print(f'Github search API returned 200 but no items found for commit {commit_sha}. Trying again...')
        else:
            print(f'GitHub search API returned {resp.status_code}: {resp.text} for commit {commit_sha}')

        retry_counter += 1

        if "You have exceeded a secondary rate limit." in resp.json().get("message", ""):
            print('Secondary rate limit exceeded. Retrying in 60 seconds...')
            time.sleep(60)

    time.sleep(2)
    return ""

def get_repo_from_task_definition(ecs, task_definition_arn, service_name=None):
    td = ecs.describe_task_definition(taskDefinition=task_definition_arn)["taskDefinition"]
    containers = td.get("containerDefinitions", [])
    for container in containers:
        image = container.get("image", "")
        if ":" in image:
            _, tag = image.rsplit(":", 1)
            if len(tag) == 40 and all(c in "0123456789abcdef" for c in tag.lower()):
                repo = get_repo_from_commit(tag)
                if repo:
                    print(f'Determined repo for service {service_name} from image tag commit {tag}. Repo: {repo}')
                    return repo
    if service_name:
        repo_candidate = service_name
        url = f"https://api.github.com/repos/{GITHUB_ORG}/{repo_candidate}"
        headers = {"Authorization": f"token {GITHUB_TOKEN}"}
        resp = requests.get(url, headers=headers)
        if resp.status_code == 200:
            repo = f"{GITHUB_ORG}/{repo_candidate}"
            print(f'Determined repo for service {service_name} from service name. Repo: {repo}')
            return repo

    print(f'Failed to determine repo for service {service_name}')
    return ""

def has_jenkinsfile(repo_full_name, path=""):
    if not repo_full_name:
        return False
    url = f"https://api.github.com/repos/{repo_full_name}/contents/{path}" if path else f"https://api.github.com/repos/{repo_full_name}/contents/"
    headers = {"Authorization": f"token {GITHUB_TOKEN}"}
    resp = requests.get(url, headers=headers)
    if resp.status_code != 200:
        return False
    files = resp.json()
    if isinstance(files, dict) and files.get("type") == "file":
        return files.get("name") == "Jenkinsfile"
    return any(f.get("name") == "Jenkinsfile" for f in files if isinstance(f, dict))

def is_in_jenkins_org_folder(repo_name):
    try:
        auth = (JENKINS_USER, JENKINS_API_TOKEN) if JENKINS_USER and JENKINS_API_TOKEN else None
        resp = requests.get(JENKINS_ORG_URL, auth=auth)
        if resp.status_code != 200:
            return False
        jobs = resp.json().get("jobs", [])
        return any(job.get("name") == repo_name for job in jobs)
    except Exception:
        return False

def gdb_ecs_service_name(service_name: str) -> str:
    for env in ENVIRONMENTS:
        if service_name.startswith(env):
            return service_name.removeprefix(f"{env}-")
    return service_name

def env_mapping(env):
    if env:
        match env:
            case "Production":
                return "prod"
            case "Staging":
                return "stage"
            case "Development":
                return "dev"
            case _:
                return env.lower()
    else:
        return "unknown"

def set_app_family(service_name):
    if service_name.startswith("apollo-"):
        return "apollo"
    if service_name.startswith("delphi-"):
        return "delphi"
    if service_name.startswith("artistapp-") or service_name.startswith("rti-"):
        return "rti"
    return ""

def list_ecs_services_with_env(ecs):
    services_info = defaultdict(lambda: {env: "" for env in ENVIRONMENTS})
    repo_map = {}
    cluster_arns = []
    paginator = ecs.get_paginator("list_clusters")
    for page in paginator.paginate():
        cluster_arns.extend(page["clusterArns"])
    for cluster_arn in cluster_arns:
        cluster_name = cluster_arn.split("/")[-1]
        service_arns = []
        svc_paginator = ecs.get_paginator("list_services")
        for svc_page in svc_paginator.paginate(cluster=cluster_arn):
            service_arns.extend(svc_page["serviceArns"])
        if not service_arns:
            continue
        for i in range(0, len(service_arns), 10):
            batch = service_arns[i:i+10]
            services = describe_services_with_backoff(ecs, cluster_arn, batch)
            for svc in services:
                tags = ecs.list_tags_for_resource(resourceArn=svc["serviceArn"]).get("tags", [])
                service_name = next((t.get("value") for t in tags if t.get("key") == "service_name"), gdb_ecs_service_name(svc["serviceName"]))
                app_family = next((t.get("value") for t in tags if t.get("key") == "application_family"), set_app_family(service_name))
                environment = next((t.get("value") for t in tags if t.get("key") == "environment"), None)

                # Always add to services_info, even if tags are missing
                key = (app_family, service_name)
                # If environment tag is present and valid, set it, else leave all envs blank
                if env_mapping(environment) in ENVIRONMENTS:
                    env_key = env_mapping(environment)
                    services_info[key][env_key] = cluster_name
                # Store repo map for all services
                if key not in repo_map or not repo_map[key]:
                    if "taskDefinition" in svc:
                        td_arn = svc["taskDefinition"]
                        repo_full_name = get_repo_from_task_definition(ecs, td_arn, service_name)
                        repo_map[key] = repo_full_name
    return services_info, repo_map

def strip_env_prefix(cluster_name):
    if not cluster_name:
        return ""
    parts = cluster_name.split("-", 1)
    return parts[1] if len(parts) == 2 else cluster_name

def describe_services_with_backoff(ecs, cluster_arn, batch):
    import time
    from botocore.exceptions import ClientError
    retries = 0
    while True:
        try:
            return ecs.describe_services(cluster=cluster_arn, services=batch)["services"]
        except ClientError as e:
            if e.response['Error']['Code'] == 'ThrottlingException':
                wait = min(1 ** retries, 30)
                print(f"ECS API throttled. Sleeping for {wait} seconds (retry {retries+1})...")
                time.sleep(wait)
                retries += 1
            else:
                raise

def skip_service(env_accounts):
    """Returns true if envs are either dev or unknown and account id is dev account, or if service only exists in dev."""
    envs = [env for env in env_accounts.keys() if len(env_accounts[env]) > 0]

    if set(envs).issubset({"dev", "unknown"}):
        account_ids = set([id for ids in env_accounts.values() for id in ids])
        if account_ids == {"103233932089"}:
            return True

    if set(envs) == {"dev"}:
        return True

    return False

def is_in_software_catalog(service_name):
    with ApiClient(dd_config) as api_client:
        api_instance = SoftwareCatalogApi(api_client)
        response = api_instance.list_catalog_entity(filter_name=service_name)
        for entity in response["data"]:
            attributes = entity.get("attributes", {})
            if attributes.get("kind") == "service" and 'metadata_origin:datadog-tools' in attributes.get("tags", []):
                return True

    return False


def main():
    current_user = iam_client.get_user()
    username = current_user['User']['UserName']
    services_info = defaultdict(lambda: {env: [] for env in ENVIRONMENTS})
    repo_map = {}

    for account in get_aws_accounts():
        print(f"Processing AWS account: {account}")
        try:
            sts = get_client(username, account, "sts")
            ecs = get_client(username, account, "ecs")
            aws_account_id = sts.get_caller_identity()["Account"]
            ecs_services, repo_lookup = list_ecs_services_with_env(ecs)
            for (app_family, service_name), env_clusters in ecs_services.items():
                raw_cluster_name = next((env_clusters[env] for env in ENVIRONMENTS if env_clusters.get(env)), "")
                logical_name = strip_env_prefix(raw_cluster_name)
                key = (app_family, service_name, logical_name)
                for env in ENVIRONMENTS:
                    if env_clusters.get(env):
                        if aws_account_id not in services_info[key][env]:
                            services_info[key][env].append(aws_account_id)
                if key not in repo_map or not repo_map[key]:
                    repo_full_name = repo_lookup.get((app_family, service_name), "")
                    if not repo_full_name:
                        repo_full_name = REPO_OVERRIDES.get(service_name, "")
                        if repo_full_name:
                            print(f"Using manual repo override for {service_name}: {repo_full_name}")
                    repo_map[key] = repo_full_name
        except Exception as e:
            print(f"Failed for account {account}: {e}")
            continue

    with open(OUTPUT_CSV, "w", newline="") as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow([
            "application_family",
            "github_repo",
            "service_name",
            "ecs_cluster_name",
            "dev",
            "qa",
            "stage",
            "uat",
            "prod",
            "shared",
            "backup",
            "networking",
            "unknown",
            "has_jenkinsfile",
            "in_jenkins_org_folder",
            "in_software_catalog",
        ])
        for (app_family, service_name, logical_name), env_accounts in services_info.items():
            if skip_service(env_accounts):
                print(f"skipping {(app_family, service_name, logical_name)}")
            else:
                repo_full_name = repo_map.get((app_family, service_name, logical_name), "")
                repo_name = repo_full_name.split("/")[-1] if repo_full_name else ""
                has_jenkins = has_jenkinsfile(repo_full_name)
                in_jenkins_org = is_in_jenkins_org_folder(repo_name)
                in_software_catalog = is_in_software_catalog(service_name)
                row = [
                    app_family or "",
                    repo_full_name or "",
                    service_name or "",
                    logical_name or ""
                ]
                for env in ENVIRONMENTS:
                    account_ids = env_accounts.get(env, [])
                    row.append(";".join(account_ids) if account_ids else "")
                row.extend([str(has_jenkins), str(in_jenkins_org), str(in_software_catalog)])
                writer.writerow(row)
    print(f"Done! CSV written to {OUTPUT_CSV}")

if __name__ == "__main__":
    main()