"""Read the already-public feature atlas and lab evidence files.

Records keep the fields stored in those files. Stage D metric rows stay labeled
with their own protocol and are not merged into the feature atlas.
"""
from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path

ROOT = Path(__file__).resolve().parent
SITE = ROOT / "site"
FEATURES = SITE / "assets" / "features.json"
EVIDENCE = SITE / "lab-experiments" / "data" / "evidence.json"

PRIVATE_TEXT = ("prompt", "original_top_tokens", "edited_top_tokens", "title")


@lru_cache(maxsize=1)
def feature_atlas() -> dict:
    document = json.loads(FEATURES.read_text(encoding="utf-8"))
    if document.get("schema") != "opensae.feature_atlas.observed_last_token.v1":
        raise ValueError("features.json is not the published last-token atlas")
    return document


@lru_cache(maxsize=1)
def lab_evidence() -> dict:
    document = json.loads(EVIDENCE.read_text(encoding="utf-8"))
    if document.get("schema") != "opensae.lab.evidence.v1":
        raise ValueError("evidence.json is not the published lab evidence file")
    return document


def _page(rows: list, limit: int, offset: int) -> dict:
    return {"total": len(rows), "offset": offset, "limit": limit, "records": rows[offset:offset + limit]}


def atlas_summary() -> dict:
    document = feature_atlas()
    features = document["features"]
    examples = [example for feature in features for example in feature.get("examples", [])]
    prompts = document["prompts"]
    return {
        "source": "assets/features.json",
        "schema": document["schema"],
        "coverage_note": document.get("coverage_note"),
        "prompt_hash_algorithm": document.get("prompt_hash_algorithm"),
        "feature_count": len(features),
        "ablation_count": len(examples),
        "prompt_count": len(prompts),
        "prompts": prompts,
        "models": sorted({feature.get("model") for feature in features}),
        "inference": "not_run",
    }


def list_features(model: str = "", q: str = "", limit: int = 25, offset: int = 0) -> dict:
    rows = []
    for feature in feature_atlas()["features"]:
        if model and feature.get("model") != model:
            continue
        haystack = " ".join(str(feature.get(key, "")) for key in ("id", "site_id", "candidate_label"))
        if q and q.lower() not in haystack.lower():
            continue
        rows.append({
            "id": feature.get("id"),
            "site_id": feature.get("site_id"),
            "feature_id": feature.get("feature_id"),
            "model": feature.get("model"),
            "layer": feature.get("layer"),
            "observed_count": feature.get("observed_count"),
            "explanation_status": feature.get("explanation_status"),
            "source_url": feature.get("source_url"),
            "schema": feature_atlas()["schema"],
        })
    return _page(rows, limit, offset)


def get_feature(feature_id: str) -> dict | None:
    for feature in feature_atlas()["features"]:
        if feature.get("id") == feature_id:
            return {"source": "assets/features.json", "schema": feature_atlas()["schema"], "record": feature, "inference": "not_run"}
    return None


def _visible_effect(record: dict) -> dict:
    copied = dict(record)
    if record.get("public_text_reviewed") is not True or record.get("text_stored") is not True:
        for key in PRIVATE_TEXT:
            if key in copied:
                copied[key] = None
        copied["text_withheld"] = "This effect is stored, but its prompt is not marked as reviewed public text."
    copied["source"] = "lab-experiments/data/evidence.json"
    copied["schema"] = lab_evidence()["schema"]
    return copied


def list_lab_records(kind: str = "effects", limit: int = 25, offset: int = 0) -> dict:
    document = lab_evidence()
    if kind not in ("effects", "coalitions", "contexts", "rss", "sites", "metrics"):
        raise ValueError("kind must be effects, coalitions, contexts, rss, sites, or metrics")
    rows = []
    for record in document[kind]:
        copied = _visible_effect(record) if kind == "effects" else dict(record)
        copied["source"] = "lab-experiments/data/evidence.json"
        copied["schema"] = document["schema"]
        if "protocol_id" not in copied and kind == "sites":
            copied["protocol_note"] = "Site identity only. Metric rows keep stage_d.historical_64_group.v1 and are not part of the feature atlas."
        rows.append(copied)
    return {
        **_page(rows, limit, offset),
        "kind": kind,
        "evidence_schema": document["schema"],
        "evidence_meta": document.get("meta"),
        "separated_from_feature_atlas": True,
        "inference": "not_run",
    }