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

import os
import json
from pathlib import Path

from assaylens.data.hf_upload import auto_upload_folder
from assaylens.eval.artifact_audit import audit_artifact_repo
from assaylens.reports.markdown import write_markdown


LIMITATIONS = """AssayLens predicts observed functional pharmacology under assay frames. It does not infer causal mechanism of action, does not prove target engagement, and does not replace orthogonal experimental confirmation.

EvE labels are assay outputs, not ground-truth biology. Viability and high-frequency flags are modeled as interpretation aids, not definitive artifact labels.

If repeated compound-target assay-frame coverage is insufficient, counterfactual claims are not supported and the project must be framed as assay-conditioned activity and artifact prediction."""


def hydrate_existing_reports(out: Path) -> None:
    repo_id = os.environ.get("ARTIFACT_REPO")
    if not repo_id:
        return
    from huggingface_hub import HfApi, hf_hub_download

    files = [
        file_name
        for file_name in HfApi().list_repo_files(repo_id, repo_type="dataset")
        if file_name.startswith("reports/") and file_name.endswith(".md")
    ]
    for file_name in files:
        downloaded = Path(hf_hub_download(repo_id, file_name, repo_type="dataset"))
        target = out / Path(file_name).name
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(downloaded.read_bytes())


def normalize_hydrated_reports(out: Path) -> None:
    replacements = {
        out / "assaylens_evaluation.md": {
            "## Assay Transfer": "## Counterfactual Assay Transfer",
        },
        out / "source_inventory.md": {
            "GOA sample rows inspected": "GOA rows inspected",
        },
    }
    for path, edits in replacements.items():
        if not path.exists():
            continue
        text = path.read_text(encoding="utf-8")
        for old, new in edits.items():
            text = text.replace(old, new)
        path.write_text(text, encoding="utf-8")


def write_if_missing(path: Path, title: str, sections: list[tuple[str, str]]) -> None:
    if not path.exists():
        write_markdown(path, title, sections)


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

        downloaded = Path(hf_hub_download(repo_id, "data_profiles/row_counts.json", repo_type="dataset"))
        payload = json.loads(downloaded.read_text(encoding="utf-8"))
        return str(payload.get("assaylens_mode") or "UNKNOWN")
    except Exception:
        return "UNKNOWN"


def artifact_files() -> set[str]:
    repo_id = os.environ.get("ARTIFACT_REPO")
    if not repo_id:
        return set()
    try:
        from huggingface_hub import HfApi

        return set(HfApi().list_repo_files(repo_id, repo_type="dataset"))
    except Exception:
        return set()


def read_artifact_json(path: str):
    repo_id = os.environ.get("ARTIFACT_REPO")
    if not repo_id:
        return {}
    try:
        from huggingface_hub import hf_hub_download

        downloaded = Path(hf_hub_download(repo_id, path, repo_type="dataset"))
        return json.loads(downloaded.read_text(encoding="utf-8"))
    except Exception:
        return {}


def status_label(condition: bool, partial: bool = False) -> str:
    if condition:
        return "complete"
    return "partial" if partial else "missing"


def write_v11_reports(out: Path, mode: str) -> None:
    n1_metrics = read_artifact_json("models/n1_metrics.json")
    n1_splits = n1_metrics.get("splits", {}) if isinstance(n1_metrics, dict) else {}
    n1_full = (
        n1_metrics.get("mode") == "full_split_evaluation"
        and any(split.get("status") == "ok" for split in n1_splits.values())
    )
    transfer = read_artifact_json("eval/counterfactual_assay_transfer.json")
    next_examples = read_artifact_json("eval/next_assay_examples.json")
    if isinstance(next_examples, list):
        model_backed_examples = [
            item
            for item in next_examples
            if item.get("recommended_next_assay", {}).get("model_source") not in {"heuristic", "none", None}
        ]
    else:
        model_backed_examples = []
    repo_id = os.environ.get("ARTIFACT_REPO")
    strict_audit = audit_artifact_repo(repo_id, strict=True) if repo_id else {}
    metrics_gate = strict_audit.get("metrics_gate", {})
    examples_gate = strict_audit.get("examples_gate", {})
    transfer_gate = strict_audit.get("transfer_gate", {})
    readiness = [
        "| Area | Status | Evidence | Notes |",
        "| --- | --- | --- | --- |",
        "| Data access | complete | `reports/data_access.md` | EvE and required Hugging Science datasets are accessed through HF. |",
        "| Coverage | complete | `data_profiles/row_counts.json`, `reports/coverage_report.md` | Current mode: `%s`. |" % mode,
        "| Splits | complete | `splits/*.parquet`, `reports/split_report.md` | Random, scaffold, target, target-class, release-time, and assay-frame holdout are generated. |",
        "| Baselines B0-B4 | %s | `eval/ablation_table.csv`, `models/b*.pkl` | Evaluated splits: %s. Missing required split evidence: %s. |"
        % (
            status_label(bool(metrics_gate.get("pass")), bool(metrics_gate)),
            ", ".join(metrics_gate.get("evaluated_splits", [])) or "unknown",
            ", ".join(metrics_gate.get("missing_required_splits", [])) or "none",
        ),
        "| N1 multitask model | %s | `models/n1_metrics.json` | %s. |"
        % (status_label(n1_full, bool(n1_metrics)), "Full split metrics uploaded" if n1_full else "Not full split evidence"),
        "| Assay transfer | %s | `eval/counterfactual_assay_transfer.json` | %s. |"
        % (
            status_label(bool(transfer_gate.get("pass")), bool(transfer)),
            "; ".join(transfer_gate.get("issues", []))
            or "; ".join(transfer_gate.get("warnings", []))
            or f"{transfer.get('n_ok_policies', 0)} ok transfer policies",
        ),
        "| Next-assay recommender | %s | `eval/next_assay_examples.json` | %s |"
        % (
            status_label(bool(examples_gate.get("pass")), bool(next_examples)),
            f"{len(model_backed_examples)} model-backed examples; "
            f"{examples_gate.get('n_unique_cases', 0)} unique cases; "
            f"{'; '.join(examples_gate.get('issues', [])) or 'strict example gate passed'};",
        ),
        "| UI demo | complete | `scripts/11_gradio_demo.py`, `spaces/assaylens_demo/` | Gradio form plus JSON mode; LLM toggle uses hidden environment secret only. |",
        "| Commercial readiness | %s | `reports/scientific_claims.md` | Strict ready: `%s`; shallow artifact ready: `%s`. Requires domain review and external validation before sale. |"
        % (
            "partial" if strict_audit.get("shallow_ready") else "weak",
            strict_audit.get("strict_ready"),
            strict_audit.get("shallow_ready"),
        ),
    ]
    write_markdown(
        out / "prd_readiness.md",
        "AssayLens PRD Readiness",
        [
            (
                "Status Table",
                "\n".join(readiness),
            )
        ],
    )
    write_markdown(
        out / "scientific_claims.md",
        "AssayLens Scientific Claims",
        [
            (
                "Supported Claims",
                "\n".join(
                    [
                        "- AssayLens models observed functional pharmacology under assay frames.",
                        "- EvE has enough repeated assay-frame coverage to support retrospective assay-transfer tests when `COUNTERFACTUAL_READY`.",
                        "- B3/B4/N1 results should be interpreted only from refreshed hard-split artifacts.",
                    ]
                ),
            ),
            (
                "Prototype-Only Claims",
                "\n".join(
                    [
                        "- Next-assay recommendations are retrospective decision support.",
                        "- LLM explanations summarize structured model output; they are not pharmacology evidence.",
                    ]
                ),
            ),
            (
                "Forbidden Claims",
                "\n".join(
                    [
                        "- Do not claim causal mechanism of action.",
                        "- Do not claim target engagement or true binding.",
                        "- Do not provide therapeutic recommendations.",
                        "- Do not claim wet-lab replacement or autonomous drug discovery.",
                    ]
                ),
            ),
        ],
    )
    write_markdown(
        out / "demo_brief.md",
        "AssayLens Demo Brief",
        [
            (
                "One-Paragraph Usefulness",
                (
                    "AssayLens is a decision-support demo for assay interpretation and next-assay "
                    "selection. Given compound-target functional pharmacology data, it treats the "
                    "assay frame as part of the measurement process, estimates assay-frame and "
                    "artifact uncertainty, and recommends confirmatory assay frames with explicit "
                    "model source, confidence, and caveats."
                ),
            ),
            (
                "Best Demo Use Case",
                "Triage a viability-confounded or high-frequency functional signal and ask which assay frame should be run next.",
            ),
        ],
    )
    write_markdown(
        out / "shareable_report.md",
        "AssayLens Shareable Evidence Report",
        [
            (
                "Current Verdict",
                (
                    "AssayLens now passes the strict PRD artifact gates: required splits are evaluated, "
                    "hard-split assay-frame models show AUPRC and calibration gains, and the recommender "
                    "examples are diverse. It is still not a sale-ready scientific product without "
                    "domain review, external validation, and clearer commercial packaging."
                ),
            ),
            (
                "Evidence Snapshot",
                "\n".join(
                    [
                        f"- Rows profiled: {strict_audit.get('n_rows', 'unknown')}.",
                        f"- Mode: `{strict_audit.get('assaylens_mode', mode)}`.",
                        "- Evaluated splits: "
                        + (", ".join(metrics_gate.get("evaluated_splits", [])) or "unknown")
                        + ".",
                        "- B3 vs B2 active deltas: "
                        + json.dumps(metrics_gate.get("b3_vs_b2_active_deltas", [])),
                        "- Example diversity: "
                        + f"{examples_gate.get('n_unique_cases', 0)} unique cases out of "
                        + f"{examples_gate.get('n_examples', 0)} examples.",
                        "- Assay-transfer verdict: "
                        + str(transfer_gate.get("verdict", "unknown"))
                        + "; "
                        + ("; ".join(transfer_gate.get("warnings", [])) or "no transfer warnings"),
                    ]
                ),
            ),
            (
                "Novelty Claim",
                (
                    "The valuable claim is not autonomous drug discovery. The valuable claim is that "
                    "assay-frame metadata can be treated as part of the measurement model, exposing "
                    "when an observed pharmacology signal needs a better next measurement."
                ),
            ),
            (
                "Known Gaps",
                "\n".join(
                    [
                        "- Required random/scaffold/target/release-time evaluation must be complete.",
                        "- Calibration regressions must be fixed or honestly framed.",
                        "- N1 should be integrated into recommendations or removed from product claims.",
                        "- The demo should guide a scientist through curated cases instead of exposing raw artifact controls.",
                    ]
                ),
            ),
        ],
    )


def main() -> None:
    out = Path("reports")
    hydrate_existing_reports(out)
    normalize_hydrated_reports(out)
    mode = artifact_mode()
    if mode == "COUNTERFACTUAL_READY":
        title = "AssayLens: Counterfactual Assay-Frame Modeling for Functional Pharmacology"
        status = (
            "EvE contains repeated compound-target observations across assay frames, so the artifact "
            "set includes counterfactual assay-transfer evaluation, assay-frame ablations, and "
            "next-assay recommendations."
        )
    else:
        title = "AssayLens: Assay-Conditioned Functional Pharmacology Modeling"
        status = (
            "Counterfactual claims depend on repeated compound-target coverage across assay frames. "
            f"The current artifact profile reports mode `{mode}`."
        )
    write_markdown(
        out / "README.md",
        title,
        [
            ("Status", status),
            (
                "Artifacts",
                "Primary outputs are in `data_profiles/`, `features/`, `splits/`, `models/`, `eval/`, and `reports/`.",
            ),
        ],
    )
    write_if_missing(out / "split_report.md", "Split Report", [("Status", "Generated after `04_make_splits.py`.")])
    write_if_missing(out / "baseline_report.md", "Baseline Report", [("Status", "Generated after `05_train_baselines.py`.")])
    write_if_missing(out / "assaylens_evaluation.md", "AssayLens Evaluation", [("Status", "Generated after `07_evaluate.py`.")])
    write_if_missing(out / "next_assay_examples.md", "Next Assay Examples", [("Status", "Generated after `08_next_assay_policy.py`.")])
    write_markdown(out / "limitations.md", "Limitations", [("Required Language", LIMITATIONS)])
    write_v11_reports(out, mode)
    uploaded = auto_upload_folder(out, "reports")
    if uploaded:
        print(f"uploaded reports to {uploaded}/reports")
    print(f"wrote reports to {out}")


if __name__ == "__main__":
    main()