#!/usr/bin/env python
from __future__ import annotations

import json
import os
import argparse
import ast
import sys
from pathlib import Path
from typing import Any

import pandas as pd

from assaylens.data.access import iter_limited_rows
from assaylens.agent.next_assay import recommend_next_assay
from assaylens.agent.model_predictor import predictor_from_artifacts
from assaylens.data.schemas import EVE_HF_ID
from assaylens.data.hf_upload import auto_upload_folder
from assaylens.reports.markdown import write_markdown


def assaylens_mode_from_artifacts() -> str:
    repo_id = os.environ.get("ARTIFACT_REPO")
    if not repo_id:
        return "ASSAY_CONDITIONED_ONLY"
    try:
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(repo_id, "data_profiles/row_counts.json", repo_type="dataset")
        return str(json.loads(Path(path).read_text(encoding="utf-8")).get("assaylens_mode"))
    except Exception:  # noqa: BLE001
        return "ASSAY_CONDITIONED_ONLY"


def frame_coverage_from_artifacts() -> dict[str, int] | None:
    repo_id = os.environ.get("ARTIFACT_REPO")
    if not repo_id:
        return None
    try:
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(repo_id, "data_profiles/frame_counts.csv", repo_type="dataset")
        frame = pd.read_csv(path)
        return {
            str(row["assay_frame_coarse"]): int(row["n_rows"])
            for _, row in frame.iterrows()
        }
    except Exception as exc:  # noqa: BLE001
        print(f"[next-assay] could not load frame coverage: {type(exc).__name__}: {exc}", file=sys.stderr)
        return None


def parse_coverage_key(key: str) -> tuple[Any, Any, Any] | None:
    try:
        value = ast.literal_eval(key)
    except (SyntaxError, ValueError):
        return None
    if not isinstance(value, tuple) or len(value) != 3:
        return None
    return value


def candidate_catalog(frame_coverage: dict[str, int] | None) -> dict[str, list[dict[str, Any]]]:
    if not frame_coverage:
        return {}
    by_target_class: dict[str, list[tuple[int, dict[str, Any]]]] = {}
    for raw_key, count in frame_coverage.items():
        parsed = parse_coverage_key(raw_key)
        if parsed is None:
            continue
        target_class, mode, mechanism = parsed
        candidate = {
            "mode": mode,
            "assay__mechanism": mechanism,
            "assay__technology": "Unknown",
        }
        by_target_class.setdefault(str(target_class), []).append((count, candidate))
    return {
        target_class: [candidate for _count, candidate in sorted(rows, key=lambda item: item[0], reverse=True)]
        for target_class, rows in by_target_class.items()
    }


def request_from_row(row: dict[str, Any], catalog: dict[str, list[dict[str, Any]]]) -> dict[str, Any] | None:
    target_class = str(row.get("target__class") or "")
    candidates = []
    observed_frame = {
        "mode": row.get("mode"),
        "assay__mechanism": row.get("assay__mechanism"),
        "assay__technology": row.get("assay__technology"),
    }
    for candidate in catalog.get(target_class, []):
        if candidate["mode"] == observed_frame["mode"] and candidate["assay__mechanism"] == observed_frame["assay__mechanism"]:
            continue
        if candidate not in candidates:
            candidates.append(candidate)
        if len(candidates) >= 5:
            break
    if len(candidates) < 2:
        return None
    return {
        "compound": {
            "compound_id": row.get("compound_id"),
            "smiles": row.get("compound__smiles"),
            "name": row.get("compound__name"),
        },
        "target": {
            "target_id": row.get("target_id"),
            "uniprot_id": row.get("target__uniprot_id"),
            "gene": row.get("target__gene"),
            "target_class": row.get("target__class"),
        },
        "observed_result": {
            "mode": row.get("mode"),
            "assay__mechanism": row.get("assay__mechanism"),
            "assay__detailed_mechanism": row.get("assay__detailed_mechanism"),
            "assay__technology": row.get("assay__technology"),
            "outcome_is_active": bool(row.get("outcome_is_active")),
            "outcome_potency_pxc50": row.get("outcome_potency_pxc50"),
            "outcome_max_activity": row.get("outcome_max_activity"),
            "is_quantified": bool(row.get("is_quantified")),
            "frequency_flag": bool(row.get("frequency_flag")),
            "viability_flag": bool(row.get("viability_flag")),
        },
        "available_assay_frames": candidates,
    }


def diverse_requests_from_eve(
    frame_coverage: dict[str, int] | None,
    min_examples: int = 20,
    row_limit: int = 5000,
) -> list[dict[str, Any]]:
    catalog = candidate_catalog(frame_coverage)
    if not catalog:
        return []
    requests = []
    seen = set()
    for row in iter_limited_rows(EVE_HF_ID, "train", row_limit):
        request = request_from_row(row, catalog)
        if request is None:
            continue
        key = (
            request["compound"].get("compound_id"),
            request["target"].get("target_id"),
            request["observed_result"].get("mode"),
            request["observed_result"].get("assay__mechanism"),
            request["observed_result"].get("assay__technology"),
        )
        if key in seen:
            continue
        seen.add(key)
        requests.append(request)
        if len(requests) >= min_examples:
            break
    return requests


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="examples/tiny_observations.json")
    parser.add_argument("--model-artifact", default="models/b3_ligand_target_assay.pkl")
    parser.add_argument("--include-llm-explanation", action="store_true")
    args = parser.parse_args()

    request_path = Path(args.input)
    mode = assaylens_mode_from_artifacts()
    predictor = predictor_from_artifacts(preferred_filename=args.model_artifact)
    frame_coverage = frame_coverage_from_artifacts()
    requests = json.loads(request_path.read_text(encoding="utf-8"))
    if isinstance(requests, dict):
        requests = [requests]
    if len(requests) < 20:
        generated = diverse_requests_from_eve(frame_coverage, min_examples=20)
        if len(generated) >= 20:
            requests = generated
    examples = [
        recommend_next_assay(
            request,
            assaylens_mode=mode,
            frame_coverage=frame_coverage,
            predictor=predictor,
            include_llm_explanation=args.include_llm_explanation,
        )
        for request in requests
    ]
    out = Path("/tmp/assaylens_next_assay_outputs/eval")
    out.mkdir(parents=True, exist_ok=True)
    (out / "next_assay_examples.json").write_text(json.dumps(examples, indent=2), encoding="utf-8")
    report_lines = []
    for example in examples[:20]:
        recommendation = example.get("recommended_next_assay", {})
        frame = recommendation.get("frame", {})
        priority = recommendation.get("priority")
        confidence = recommendation.get("confidence")
        model_source = recommendation.get("model_source")
        report_lines.append(
            f"- {example['compound_id']} / {example['target_id']}: "
            f"{frame.get('mode', 'unknown')} {frame.get('assay__mechanism', 'unknown')} "
            f"(priority={priority}, confidence={confidence}, model={model_source})"
        )
    reports_dir = Path("/tmp/assaylens_next_assay_outputs/reports")
    write_markdown(
        reports_dir / "next_assay_examples.md",
        "Next Assay Examples",
        [("Examples", "\n".join(report_lines))],
    )
    uploaded = auto_upload_folder(out, "eval")
    uploaded_reports = auto_upload_folder(reports_dir, "reports")
    if uploaded:
        print(f"uploaded next-assay examples to {uploaded}/eval")
    if uploaded_reports:
        print(f"uploaded next-assay report to {uploaded_reports}/reports")
    print(f"wrote {out / 'next_assay_examples.json'}")


if __name__ == "__main__":
    main()