# afusion/api.py

import os
import json
import re
import uuid
import pandas as pd
from afusion.execution import run_alphafold
from afusion.utils import compress_output_folder
from loguru import logger


def create_batch_task(job_name, entities, model_seeds, bonded_atom_pairs=None, user_ccd=None):
    """
    Creates a batch task dictionary for a single prediction.

    :param job_name: Name of the job.
    :type job_name: str
    :param entities: List of dictionaries, each representing an Entity.
        Each entity dict should have keys:
        - 'type': 'protein', 'rna', 'dna', or 'ligand'
        - 'id': str or list
        - 'sequence_data': dict with sequence information
    :type entities: list
    :param model_seeds: List of integers.
    :type model_seeds: list of int
    :param bonded_atom_pairs: Optional list of bonded atom pairs.
    :type bonded_atom_pairs: list, optional
    :param user_ccd: Optional user CCD.
    :type user_ccd: str, optional
    :return: Dictionary representing the AlphaFold input JSON structure.
    :rtype: dict
    """
    sequences = []
    for entity in entities:
        entity_type = entity['type']
        sequence_data = entity['sequence_data']
        entity_id = entity['id']

        sequence_entry = sequence_data.copy()
        sequence_entry['id'] = entity_id

        if entity_type == 'protein':
            sequences.append({'protein': sequence_entry})
        elif entity_type == 'rna':
            sequences.append({'rna': sequence_entry})
        elif entity_type == 'dna':
            sequences.append({'dna': sequence_entry})
        elif entity_type == 'ligand':
            sequences.append({'ligand': sequence_entry})
        else:
            logger.error(f"Unknown entity type: {entity_type}")
            continue

    alphafold_input = {
        "name": job_name,
        "modelSeeds": model_seeds,
        "sequences": sequences,
        "dialect": "alphafold3",
        "version": 1
    }

    if bonded_atom_pairs:
        alphafold_input["bondedAtomPairs"] = bonded_atom_pairs

    if user_ccd:
        alphafold_input["userCCD"] = user_ccd

    logger.debug(f"Created task for job: {job_name}")
    return alphafold_input


def run_batch_predictions(
    tasks,
    af_input_base_path,
    af_output_base_path,
    model_parameters_dir,
    databases_dir,
    run_data_pipeline=True,
    run_inference=True,
    bucket_sizes=None,
):
    """
    Runs batch predictions for the given tasks.

    :param tasks: List of task dicts, as generated by create_batch_task.
    :type tasks: list of dict
    :param af_input_base_path: Base path for AlphaFold input.
    :type af_input_base_path: str
    :param af_output_base_path: Base path for AlphaFold output.
    :type af_output_base_path: str
    :param model_parameters_dir: Path to model parameters directory.
    :type model_parameters_dir: str
    :param databases_dir: Path to databases directory.
    :type databases_dir: str
    :param run_data_pipeline: Whether to run data pipeline.
    :type run_data_pipeline: bool
    :param run_inference: Whether to run inference.
    :type run_inference: bool
    :param bucket_sizes: Optional list of bucket sizes.
    :type bucket_sizes: list of int, optional
    :return: List of dicts with keys 'job_name', 'output_folder', 'status'.
    :rtype: list of dict
    """
    results = []
    for task in tasks:
        job_name = task['name']
        job_folder_name = job_name
        
        input_path = os.path.join(af_input_base_path, job_folder_name)
        output_path = af_output_base_path
        
        os.makedirs(input_path, exist_ok=True)
        os.makedirs(output_path, exist_ok=True) 

        json_save_path = os.path.join(input_path, "fold_input.json")
        try:
            with open(json_save_path, "w") as json_file:
                json.dump(task, json_file, indent=2)
            logger.info(f"JSON file saved for job '{job_name}' at {json_save_path}")
        except Exception as e:
            logger.error(f"Error saving JSON file for job '{job_name}': {e}")
            results.append({
                'job_name': job_name,
                'output_folder': output_path,
                'status': f'Failed to save JSON: {e}'
            })
            continue

        # Build the Docker command
        docker_command = (
            f"docker run --rm "
            f"--volume {input_path}:/root/af_input "
            f"--volume {output_path}:/root/af_output "
            f"--volume {model_parameters_dir}:/root/models "
            f"--volume {databases_dir}:/root/public_databases "
            f"--gpus all "
            f"alphafold3 "
            f"python run_alphafold.py "
            f"--json_path=/root/af_input/fold_input.json "
            f"--model_dir=/root/models "
            f"--output_dir=/root/af_output "
            f"{'--run_data_pipeline' if run_data_pipeline else ''} "
            f"{'--run_inference' if run_inference else ''} "
            f"{'--buckets ' + ','.join(map(str, bucket_sizes)) if bucket_sizes else ''}"
        )

        logger.debug(f"Running Docker command for job '{job_name}': {docker_command}")

        # Run AlphaFold
        try:
            output = run_alphafold(docker_command)
            logger.info(f"AlphaFold execution completed for job '{job_name}'.")

            # Check if the output directory exists
            expected_output_folder = os.path.join(af_output_base_path, job_folder_name)
            if os.path.exists(expected_output_folder):
                logger.info(f"Results saved in: {expected_output_folder}")
                status = 'Success'
            else:
                logger.error(f"Output folder '{expected_output_folder}' not found for job '{job_name}'.")
                status = 'Failed'
        except Exception as e:
            logger.error(f"Error running AlphaFold for job '{job_name}': {e}")
            status = f'Failed to run AlphaFold: {e}'

        results.append({
            'job_name': job_name,
            'output_folder': output_path,
            'status': status
        })

    return results


def create_protein_sequence_data(sequence, modifications=None, msa_option='auto', unpaired_msa=None, paired_msa=None, templates=None):
    """
    Creates sequence data for a protein entity.

    :param sequence: The protein sequence.
    :type sequence: str
    :param modifications: Optional list of modifications, each with keys 'ptmType' and 'ptmPosition'.
    :type modifications: list of dict, optional
    :param msa_option: MSA option, 'auto', 'upload', or 'none'.
    :type msa_option: str
    :param unpaired_msa: Unpaired MSA (if msa_option is 'upload').
    :type unpaired_msa: str, optional
    :param paired_msa: Paired MSA (if msa_option is 'upload').
    :type paired_msa: str, optional
    :param templates: Optional list of template dicts.
    :type templates: list of dict, optional
    :return: Sequence data dictionary.
    :rtype: dict
    """
    protein_entry = {
        "sequence": sequence
    }
    if modifications:
        protein_entry["modifications"] = modifications
    if msa_option == 'auto':
        protein_entry["unpairedMsa"] = None
        protein_entry["pairedMsa"] = None
        protein_entry["templates"] = []
    elif msa_option == 'none':
        protein_entry["unpairedMsa"] = ""
        protein_entry["pairedMsa"] = ""
        protein_entry["templates"] = []
    elif msa_option == 'upload':
        protein_entry["unpairedMsa"] = unpaired_msa or ""
        protein_entry["pairedMsa"] = paired_msa or ""
        protein_entry["templates"] = templates or []
    else:
        logger.error(f"Invalid msa_option: {msa_option}")
    return protein_entry


def create_rna_sequence_data(sequence, modifications=None, msa_option='auto', unpaired_msa=None):
    """
    Creates sequence data for an RNA entity.

    :param sequence: The RNA sequence.
    :type sequence: str
    :param modifications: Optional list of modifications.
    :type modifications: list of dict, optional
    :param msa_option: MSA option, 'auto', 'upload', or 'none'.
    :type msa_option: str
    :param unpaired_msa: Unpaired MSA (if msa_option is 'upload').
    :type unpaired_msa: str, optional
    :return: Sequence data dictionary.
    :rtype: dict
    """
    rna_entry = {
        "sequence": sequence
    }
    if modifications:
        rna_entry["modifications"] = modifications
    if msa_option == 'auto':
        rna_entry["unpairedMsa"] = None
    elif msa_option == 'none':
        rna_entry["unpairedMsa"] = ""
    elif msa_option == 'upload':
        rna_entry["unpairedMsa"] = unpaired_msa or ""
    else:
        logger.error(f"Invalid msa_option: {msa_option}")
    return rna_entry


def create_dna_sequence_data(sequence, modifications=None):
    """
    Creates sequence data for a DNA entity.

    :param sequence: The DNA sequence.
    :type sequence: str
    :param modifications: Optional list of modifications.
    :type modifications: list of dict, optional
    :return: Sequence data dictionary.
    :rtype: dict
    """
    dna_entry = {
        "sequence": sequence
    }
    if modifications:
        dna_entry["modifications"] = modifications
    return dna_entry


def create_ligand_sequence_data(ccd_codes=None, smiles=None):
    """
    Creates sequence data for a ligand entity.

    :param ccd_codes: List of CCD codes.
    :type ccd_codes: list of str, optional
    :param smiles: SMILES string.
    :type smiles: str, optional
    :return: Sequence data dictionary.
    :rtype: dict
    """
    if ccd_codes and smiles:
        logger.error("Please provide only one of CCD Codes or SMILES String.")
        return {}
    elif ccd_codes:
        ligand_entry = {
            "ccdCodes": ccd_codes
        }
        return ligand_entry
    elif smiles:
        ligand_entry = {
            "smiles": smiles
        }
        return ligand_entry
    else:
        logger.error("Ligand requires either CCD Codes or SMILES String.")
        return {}


def create_tasks_from_dataframe(df):
    """
    Creates batch tasks from a DataFrame.

    :param df: DataFrame with columns representing parameters:
        - 'job_name': str
        - 'type': 'protein', 'rna', 'dna', or 'ligand'
        - 'id': str or list
        - 'sequence': str
        - Other optional parameters:
            - 'modifications': list of dicts (as JSON string)
            - 'msa_option': 'auto', 'upload', or 'none'
            - 'unpaired_msa': str
            - 'paired_msa': str
            - 'templates': list of dicts (as JSON string)
            - 'model_seeds': list of integers (as string)
            - 'bonded_atom_pairs': list (as JSON string)
            - 'user_ccd': str
    :type df: pandas.DataFrame
    :return: List of task dictionaries.
    :rtype: list of dict
    """
    tasks = []
    grouped = df.groupby('job_name')
    for job_name, group in grouped:
        entities = []
        model_seeds = None
        bonded_atom_pairs = None
        user_ccd = None

        for _, row in group.iterrows():
            entity_type = row['type']
            entity_id = row['id']
            sequence = row.get('sequence', '')

            # Parse optional fields
            modifications = parse_json_field(row.get('modifications'))
            msa_option = row.get('msa_option', 'auto')
            unpaired_msa = row.get('unpaired_msa')
            paired_msa = row.get('paired_msa')
            templates = parse_json_field(row.get('templates'))

            # Create sequence data based on entity type
            if entity_type == 'protein':
                sequence_data = create_protein_sequence_data(
                    sequence=sequence,
                    modifications=modifications,
                    msa_option=msa_option,
                    unpaired_msa=unpaired_msa,
                    paired_msa=paired_msa,
                    templates=templates
                )
            elif entity_type == 'rna':
                sequence_data = create_rna_sequence_data(
                    sequence=sequence,
                    modifications=modifications,
                    msa_option=msa_option,
                    unpaired_msa=unpaired_msa
                )
            elif entity_type == 'dna':
                sequence_data = create_dna_sequence_data(
                    sequence=sequence,
                    modifications=modifications
                )
            elif entity_type == 'ligand':
                ccd_codes = parse_list_field(row.get('ccd_codes'))
                smiles = row.get('smiles')
                sequence_data = create_ligand_sequence_data(
                    ccd_codes=ccd_codes,
                    smiles=smiles
                )
            else:
                logger.error(f"Unknown entity type: {entity_type}")
                continue

            entities.append({
                'type': entity_type,
                'id': entity_id,
                'sequence_data': sequence_data
            })

            # Get job-level parameters (assuming they are the same for all entities in the group)
            if model_seeds is None and pd.notna(row.get('model_seeds')):
                model_seeds = parse_list_field(row.get('model_seeds'), data_type=int)
            if bonded_atom_pairs is None and pd.notna(row.get('bonded_atom_pairs')):
                bonded_atom_pairs = parse_json_field(row.get('bonded_atom_pairs'))
            if user_ccd is None and pd.notna(row.get('user_ccd')):
                user_ccd = row.get('user_ccd')

        if model_seeds is None:
            model_seeds = [1]  # Default seed if not provided

        task = create_batch_task(
            job_name=job_name,
            entities=entities,
            model_seeds=model_seeds,
            bonded_atom_pairs=bonded_atom_pairs,
            user_ccd=user_ccd
        )
        tasks.append(task)
    return tasks


def parse_json_field(value):
    """
    Parses a JSON string field into a Python object.

    Returns None if the value is NaN or empty.

    :param value: JSON string to parse.
    :type value: str
    :return: Parsed Python object or None.
    :rtype: object or None
    """
    if pd.isna(value) or value == '':
        return None
    try:
        return json.loads(value)
    except json.JSONDecodeError as e:
        logger.error(f"JSON decode error: {e}")
        return None

def parse_list_field(value, data_type=str):
    """
    Parses a comma-separated string field into a list.

    Returns None if the value is NaN or empty.

    :param value: Comma-separated string to parse.
    :type value: str
    :param data_type: Data type to convert list items to.
    :type data_type: type
    :return: List of items converted to data_type, or None.
    :rtype: list or None
    """
    if pd.isna(value) or value == '':
        return None
    return [data_type(item.strip()) for item in value.split(',') if item.strip()]

