import json
from pathlib import Path
import subprocess
from investigation import storage_path, entry_path, cwd, env
from investigation.command import git_check, check_entry, init_entry, commit_entry, run_entry, sample_entry, delete_entry, batch_submit, reset_entry, clean_runtime_cache, insert_to_entry, get_job_state, sync_src_to_commit, entry_stat, roll_script
from investigation.database.query import filter_entry
from investigation.cli.utils import auto_extract_args
from investigation.cli.args import CLI, Init, Commit, Submit, Sample, Filter, Clean, Params, Stat, Sync, Analyze, Roll
from investigation.analyze.analyze import aggregate_analysis
from slurming.Slurm.slurm import get_job_dict


def main():

    git_check()

    check_entry()
    parsed_args, unknown_args = auto_extract_args(CLI)

    # print(parsed_args)
    # input(unknown_args)

    match parsed_args:

        case Init():
            init_entry(parsed_args.name)
        case Commit():
            commit_entry(parsed_args.name, parsed_args.commit)
        case Submit():
            from investigation.entry import hostname

            cluster_name = env.str("CLUSTER_NAME")
            _, cases = filter_entry(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                **unknown_args,
            )
            # active_jobs, running_jobs, pending_jobs = get_job_dict(cluster_name=hostname)
            # for job_name in running_jobs.keys():
            #     cases.pop(job_name, None)
            # if (not parsed_args.include_pending):
            #     for job_name in pending_jobs.keys():
            #         cases.pop(job_name, None)
            job_state_dict = entry_stat(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                cluster_name=hostname,
                unknown_args=unknown_args,
                all_cluster=False,
            )
            for state in ["RUNNING", "DONE"]:
                for job_name in job_state_dict[state]:
                    cases.pop(job_name, None)
            if (not parsed_args.include_pending):
                for job_name in job_state_dict["PENDING"]:
                    cases.pop(job_name, None)
            run_entry(
                commit_hash=parsed_args.commit,
                parellels=parsed_args.parellels,
                force=parsed_args.force,
                cases=cases,
                cluster=cluster_name,
            )
        case Sample():

            from investigation.entry import hostname
            cases_df, cases = filter_entry(parsed_args.commit, batch_hash=parsed_args.batch, **unknown_args)
            print(f"Found {len(cases)} pending shell scripts.")
            if (not parsed_args.force):
                active_jobs, running_jobs, pending_jobs = get_job_dict(cluster_name=hostname)
                # print(active_jobs)
                # running_jobs = {k: v for k, v in active_jobs.items() if v['state'] in ['RUNNING']}
                # pending_jobs = {k: v for k, v in active_jobs.items() if v['state'] in ['PENDING']}

                joint_keys = set(cases.keys()).intersection(set(running_jobs.keys()))

                jobs_on_observation = {**running_jobs, **pending_jobs}

                for key in joint_keys:
                    cases.pop(key)
                # print(f"Removed running casees, {len(cases)} pending shell scripts.")

                if (not parsed_args.include_pending):
                    joint_pending_keys = set(cases.keys()).intersection(set(pending_jobs.keys()))
                    for key in joint_pending_keys:
                        cases.pop(key)
                    # print(f"Removed pending casees, {len(cases)} pending shell scripts.")
                job_states = {
                    x: get_job_state(
                        parsed_args.commit,
                        x,
                        cases[x],
                    ) if (x not in jobs_on_observation) else "PENDING"
                    for x in cases.keys()
                }
                # map(
                #     lambda x: job_states.update({
                #         x:
                #         get_job_state(parsed_args.commit, x, cases[x]) if (x not in jobs_on_observation) else "PENDING"
                #     }), cases.keys())

                for jname, jstate in job_states.items():
                    if (jstate in [
                            "UNINITIALIZED",
                            "FAILED",
                            "SAMPLED",
                            "PENDING",
                            "STARTED",
                    ]):
                        continue
                    cases.pop(jname, None)
                    # print(f"Removed {jname} with state {jstate}, {len(cases)} pending shell scripts.")
                print(f"Removed casees, {len(cases)} pending shell scripts.")

            num_resampled_cases = sample_entry(parsed_args.commit, force=parsed_args.force, cases=cases)
            print(f"Generated {num_resampled_cases} pending shell scripts.")

        case Filter():
            cases_df, cases = filter_entry(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                **unknown_args,
            )

            # print(cases_df.shape)
            case_display_df = cases_df[[
                x for x in cases_df.columns if len(cases_df[x].unique()) > 1 and x not in ["num_gpus", "num_cpus"]
            ]]
            match parsed_args.action:
                case "show":
                    # print("\n".join(cases.keys()))
                    for col in list(case_display_df.columns):
                        if len(case_display_df[col].dropna().unique()) == 1:
                            case_display_df = case_display_df.drop(columns=[col])
                    if ("name" in case_display_df.columns):
                        case_display_df.sort_values(by="name", inplace=True)
                    if (parsed_args.seed_merge):
                        case_display_df = case_display_df.groupby(
                            [col for col in case_display_df.columns if col != 'seed'], as_index=False,
                            dropna=False).agg(func='count')
                    else:
                        # 查找出name外其他列都相同的行，然后删除
                        asdf = case_display_df.reset_index()

                        duplicate_rows = asdf.duplicated(subset=[col for col in asdf.columns if col != 'job_name'],
                                                         keep=False)
                        keep_first = asdf.duplicated(subset=[col for col in asdf.columns if col != 'job_name'],
                                                     keep='first')
                        duplicate_rows = asdf[duplicate_rows & ~keep_first]
                        print(f"Found {len(duplicate_rows)} duplicate rows (excluding 'name' column).")
                        print(duplicate_rows['job_name'].to_string(index=False))
                    print("Displaying filtered cases:")
                    print(case_display_df.to_string(index=True))
                    print(case_display_df.shape)
                case "failed":
                    from investigation.entry import hostname
                    active_jobs, running_jobs, pending_jobs = get_job_dict(cluster_name=hostname)
                    for job_name in [*running_jobs.keys(), *pending_jobs.keys()]:
                        cases.pop(job_name, None)

                    for cname in list(cases.keys()):
                        case_config = cases[cname]

                        job_state = get_job_state(parsed_args.commit, cname, case_config)

                        if (job_state in ["DONE", "PENDING", "SAMPLED"]):
                            cases.pop(cname)
                            continue

                    print("\n".join(cases.keys()))
                case "delete":
                    for case_hash in cases.keys():
                        delete_entry(parsed_args.commit, case_hash, cases[case_hash])
                case "reset":
                    for case_hash in cases.keys():
                        reset_entry(parsed_args.commit, case_hash, cases[case_hash])
                case "stop":
                    from investigation.entry import hostname
                    active_jobs, running_jobs, pending_jobs = get_job_dict(cluster_name=hostname)
                    jobs_on_observation = {**running_jobs, **pending_jobs}
                    to_stop = filter(lambda x: x in jobs_on_observation, cases.keys())
                    job_ids = list(map(lambda x: str(jobs_on_observation[x]['job_id']), to_stop))
                    if (len(job_ids) > 0):
                        print(f"scancel {' '.join(job_ids)}")
                        # if len(list(job_ids)) > 0:
                        subprocess.run(["scancel", *job_ids])
                case "list":
                    for case_hash in cases.keys():
                        print(case_hash)
                case _:
                    print(
                        f"Unknown filter action {parsed_args.action}, please use 'show', 'failed', 'delete', 'reset' or 'stop'."
                    )
                    exit(1)

        case Clean():
            from investigation.entry import hostname
            case_df, cases = filter_entry(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                **unknown_args,
            )
            active_jobs, running_jobs, pending_jobs = get_job_dict(cluster_name=hostname)
            cluster_view = {**running_jobs, **pending_jobs}
            for case_hash in cases.keys():

                if case_hash in cluster_view:
                    # print(f"Case {case_hash} is still active, skipping deletion.")
                    continue
                else:
                    clean_runtime_cache(
                        parsed_args.commit,
                        case_hash,
                        cases[case_hash],
                        parsed_args.dry_run,
                    )
        case Params():
            case_df, cases = filter_entry(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                **unknown_args,
            )
            params_dict = json.loads((cwd / f"{parsed_args.add}.json").read_text())
            for case_hash in cases.keys():
                insert_to_entry(
                    parsed_args.commit,
                    case_hash,
                    cases[case_hash],
                    params_dict,
                    dry_run=parsed_args.dry_run,
                )
        case Stat():
            from investigation.entry import hostname
            job_stat_dict = entry_stat(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                cluster_name=hostname,
                unknown_args=unknown_args,
                all_cluster=True,
            )
            for state, cases in job_stat_dict.items():
                print(f"{state}: {len(cases)}")
            print(job_stat_dict["FAILED"])

        case Sync():
            sync_src_to_commit(cwd, parsed_args.commit)
        case Analyze():
            aggregate_analysis(parsed_args, unknown_args)
        case Roll():
            from investigation.entry import hostname
            job_stat_dict = entry_stat(
                parsed_args.commit,
                batch_hash=parsed_args.batch,
                unknown_args=unknown_args,
                cluster_name=hostname,
                all_cluster=False,
            )
            print({k: len(v) for k, v in job_stat_dict.items()})
            failed_cases = job_stat_dict["FAILED"]
            for case_hash in failed_cases:
                print(f"Rolling back case {case_hash}...")
                roll_script(parsed_args.commit, case_hash)

        case _:
            print(f"Unknown action {parsed_args.action}, please use 'init' or 'commit'.")
            exit(1)


if __name__ == "__main__":

    main()
