# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

from mirage.accessor.langfuse import LangfuseAccessor
from mirage.cache.index import IndexEntry
from mirage.core.hierarchy.readdir import make_readdir
from mirage.core.hierarchy.scope import ScopeMatch
from mirage.core.langfuse.client import (fetch_dataset_items,
                                         fetch_dataset_runs, fetch_datasets,
                                         fetch_prompts, fetch_sessions,
                                         fetch_traces)
from mirage.core.langfuse.scope import TOP_LEVEL_DIRS, detect_scope
from mirage.core.render.json import jsonl_bytes


async def _list_traces(accessor: LangfuseAccessor,
                       match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    traces = await fetch_traces(
        accessor.api,
        limit=accessor.config.default_trace_limit,
        from_timestamp=accessor.config.default_from_timestamp,
    )
    # The list endpoint returns trace summaries while a read renders the
    # full trace with its observations, so a size here would cost one
    # fetch_trace per entry. Traces and prompts stay size-unknown until a
    # read hydrates them; the dataset .jsonl files below are sized
    # because their listing already carries every item.
    return [(f"{t.get('id', '')}.json",
             IndexEntry(
                 id=t.get("id", ""),
                 name=t.get("id", ""),
                 resource_type="langfuse/trace",
                 vfs_name=f"{t.get('id', '')}.json",
             )) for t in traces]


async def _list_sessions(accessor: LangfuseAccessor,
                         match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    sessions = await fetch_sessions(accessor.api)
    return [(s.get("id", ""),
             IndexEntry(
                 id=s.get("id", ""),
                 name=s.get("id", ""),
                 resource_type="langfuse/session",
                 vfs_name=s.get("id", ""),
             )) for s in sessions]


async def _list_session_traces(
        accessor: LangfuseAccessor,
        match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    traces = await fetch_traces(
        accessor.api,
        session_id=match.slots["session_id"],
        limit=accessor.config.default_trace_limit,
        from_timestamp=accessor.config.default_from_timestamp,
    )
    return [(f"{t.get('id', '')}.json",
             IndexEntry(
                 id=t.get("id", ""),
                 name=t.get("id", ""),
                 resource_type="langfuse/trace",
                 vfs_name=f"{t.get('id', '')}.json",
             )) for t in traces]


async def _list_prompts(accessor: LangfuseAccessor,
                        match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    prompts = await fetch_prompts(accessor.api)
    seen: set[str] = set()
    entries: list[tuple[str, IndexEntry]] = []
    for p in prompts:
        prompt_name = p.get("name", "")
        if prompt_name in seen:
            continue
        seen.add(prompt_name)
        entries.append((prompt_name,
                        IndexEntry(
                            id=prompt_name,
                            name=prompt_name,
                            resource_type="langfuse/prompt",
                            vfs_name=prompt_name,
                        )))
    return entries


async def _list_prompt_versions(
        accessor: LangfuseAccessor,
        match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    prompt_name = match.slots["prompt_name"]
    prompts = await fetch_prompts(accessor.api)
    entries: list[tuple[str, IndexEntry]] = []
    for p in prompts:
        if p.get("name") != prompt_name:
            continue
        # The list endpoint returns PromptMeta, which carries every
        # version of a prompt in a `versions` array; there is no scalar
        # `version`.
        for version in sorted(p.get("versions", [])):
            filename = f"{version}.json"
            entries.append((filename,
                            IndexEntry(
                                id=f"{prompt_name}/{version}",
                                name=str(version),
                                resource_type="langfuse/prompt_version",
                                vfs_name=filename,
                            )))
    return entries


async def _list_datasets(accessor: LangfuseAccessor,
                         match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    datasets = await fetch_datasets(accessor.api)
    return [(d.get("name", ""),
             IndexEntry(
                 id=d.get("name", ""),
                 name=d.get("name", ""),
                 resource_type="langfuse/dataset",
                 vfs_name=d.get("name", ""),
             )) for d in datasets]


async def _list_dataset(accessor: LangfuseAccessor,
                        match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    dataset_name = match.slots["dataset_name"]
    # One dataset_items call per dataset directory actually entered: the
    # dataset listing carries no item payloads, so items.jsonl can only
    # be sized here, and only for datasets the caller opens.
    items = await fetch_dataset_items(accessor.api, dataset_name)
    return [
        ("items.jsonl",
         IndexEntry(
             id=f"{dataset_name}/items",
             name="items.jsonl",
             resource_type="langfuse/dataset_items",
             vfs_name="items.jsonl",
             size=len(jsonl_bytes(items)),
         )),
        ("runs",
         IndexEntry(
             id=f"{dataset_name}/runs",
             name="runs",
             resource_type="langfuse/dataset_runs_dir",
             vfs_name="runs",
         )),
    ]


async def _list_dataset_runs(
        accessor: LangfuseAccessor,
        match: ScopeMatch) -> list[tuple[str, IndexEntry]]:
    dataset_name = match.slots["dataset_name"]
    runs = await fetch_dataset_runs(accessor.api, dataset_name)
    entries: list[tuple[str, IndexEntry]] = []
    for r in runs:
        run_name = r.get("name", "")
        filename = f"{run_name}.jsonl"
        # The listing already carries the run document read() renders, so
        # each run file's exact size is free here.
        entries.append((filename,
                        IndexEntry(
                            id=run_name,
                            name=run_name,
                            resource_type="langfuse/dataset_run",
                            vfs_name=filename,
                            size=len(jsonl_bytes([r])),
                        )))
    return entries


readdir = make_readdir(
    detect_scope,
    listers={
        "traces": _list_traces,
        "sessions": _list_sessions,
        "session": _list_session_traces,
        "prompts": _list_prompts,
        "prompt": _list_prompt_versions,
        "datasets": _list_datasets,
        "dataset": _list_dataset,
        "runs": _list_dataset_runs,
    },
    static_root=tuple(TOP_LEVEL_DIRS),
)
