from chatbot.logic.target_search import search_protein
from chatbot.logic.bioactivity import get_bioactivity_data
from chatbot.logic.lipinski import filter_lipinski
from chatbot.logic.admet import run_admet_prediction, filter_admet_candidates
from chatbot.models.admet_model import ADMETModel
import pandas as pd

def run_pipeline(query_or_target_id, user_id, apply_lipinski=False, lipinski_threshold=4, apply_admet=False, exploratory=False):
    print("Received user query:", query_or_target_id)

    # MODE 1: Search if query is text-based
    if not (apply_lipinski or apply_admet):
        print("Step 1: Initiating target search...")
        targets = search_protein(query_or_target_id)
        print(f"Step 2: Target search completed: {len(targets)} target(s) found.")

        if not targets:
            return {"response": f"No matching targets found for ***{query_or_target_id}***."}
        if len(targets) > 1:
            print("Awaiting user selection from multiple targets.")
            return {
                "response": f"I found multiple targets matching ***{query_or_target_id}***. Please reply with the number of the target you'd like to explore, e.g. 1.",
                "targets": targets
            }

        # Single match found — proceed to prediction options
        target_raw = targets[0]
        target = {
            'id': target_raw.get("target_chembl_id", target_raw.get("id", "unknown ID")),
            'name': target_raw.get("pref_name", target_raw.get("name", "Unnamed Target")),
            'organism': target_raw.get("organism", "unknown organism"),
            'target_type': target_raw.get("target_type", "")
        }

        target_name = target["name"]
        target_organism = target["organism"]
        return {
            "response": f"You selected ***{target_name}*** from ***{target_organism}***. Would you prefer a **strict** or **lenient** prediction?",
            "options": ["Strict", "Lenient"],
            "target": target
        }

    # MODE 2: Filtering logic (user has chosen prediction mode)
    print(f"Step 4: Prediction mode selected: {'Strict' if not exploratory else 'Lenient'}")
    target_id = query_or_target_id
    try:
        target_info = search_protein(target_id)
        target_name = target_info[0]["name"] if target_info else target_id
        target_organism = target_info[0].get("organism", "unknown organism") if target_info else "unknown organism"
    except Exception:
        target_name = target_id
        target_organism = "unknown organism"

    # Retrieve bioactivity and select only active compounds
    print("Step 5: Initializing bioactivity data search...")
    bio_df = get_bioactivity_data(target_id)
    if bio_df is None or bio_df.empty:
        return {"response": f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try a different target."}

    if "bioactivity_class" not in bio_df.columns:
        return {"response": "Bioactivity data found, but no 'bioactivity_class' column is available."}
    bio_df["bioactivity_class"] = bio_df["bioactivity_class"].astype(str).str.lower()
    actives = bio_df[bio_df["bioactivity_class"] == "active"]

    print(f"Step 6: Active compounds found: {len(actives)}")
    if actives.empty:
        message = (
            f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try the **lenient** prediction or a different target!"
            if not exploratory else
            f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try a different target."
        )
        return {"response": message}

    # Lipinski filtering
    print("Step 7: Running Lipinski descriptor analysis...")
    lipinski_pass = filter_lipinski(actives, threshold=(4 if not exploratory else 3))
    if lipinski_pass is None or lipinski_pass.empty:
        message = (
            f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try the **lenient** prediction or a different target!"
            if not exploratory else
            f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try a different target."
        )
        return {"response": message}
    print(f"Step 8: Compounds passing Lipinski filter: {len(lipinski_pass)}")

    # ADMET filtering
    print("Step 9: Starting ADMET prediction analysis...")
    try:
        admet_model = ADMETModel()
        admet_df = run_admet_prediction(lipinski_pass, admet_model)
        hits = filter_admet_candidates(admet_df, strict=not exploratory)
        print(f"Step 10: Compounds satisfying ADMET properties: {len(hits)}")

        if hits.empty:
            message = (
                f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try the **lenient** prediction or a different target!"
                if not exploratory else
                f"Hmm, I couldn’t find anything promising for ***{target_name}*** from ***{target_organism}***. Try a different target."
            )
            return {"response": message}
    except Exception as e:
        print("ADMET error:", e)
        return {"response": f"ADMET prediction failed for ***{target_name}***. Try again later."}

    # Prepare output for frontend rendering
    display_cols = ["molecule_chembl_id", "canonical_smiles"]
    hits = hits[display_cols].copy()
    hits.index += 1
    hits.insert(0, "No.", hits.index)
    print("Pipeline completed and ready for frontend rendering.")

    return {
        "response": f"Here {'is' if len(hits) == 1 else 'are'} {len(hits)} compound(s) predicted to act on ***{target_name}*** from ***{target_organism}***.",
        "molecules": hits.to_dict(orient="records")
    }