from __future__ import annotations

from pydantic import BaseModel

CorpusId = str
AuthorId = str


class YearRange(BaseModel):
    start: int | None
    end: int | None


class RelevanceCriterion(BaseModel):
    name: str
    description: str
    weight: float

    @classmethod
    def from_dict(cls, data: dict) -> RelevanceCriterion:
        return RelevanceCriterion(
            name=data["name"],
            description=data["description"],
            weight=data["weight"],
        )


class BroadScorerCriteria(BaseModel):
    known_to_be_good: list[CorpusId]  # non-exhaustive, mostly for cache
    known_to_be_bad: list[
        CorpusId
    ]  # TODO - rethink this (we currently dont have this annotated in the data)
    relevance_criteria: list[RelevanceCriterion]  # prompt to llm to judge relevance
    # TODO - rethink these, how will we use them?
    venue_must_be_one_of: list[str] | None = None
    author_must_be_one_of: list[AuthorId] | None = None
    year_must_be_in_range: YearRange | None = None


class SpecificPaperScorerCriteria(BaseModel):
    corpus_ids: list[CorpusId]  # short list of Corpus IDs


ScorerCriteria = BroadScorerCriteria | SpecificPaperScorerCriteria


class Query(BaseModel):
    query_id: str
    query: str


class SingleResult(BaseModel):
    paper_id: CorpusId
    markdown_evidence: str  # should/can be markdown


class ExpectedAgentOutput(BaseModel):
    query_id: str | None = None
    results: list[SingleResult]


class AgentDataPoint(BaseModel):
    input: Query


class AgentFacingDataset(BaseModel):
    dataset: list[AgentDataPoint]


# this information should/might be hidden from users
class ScorerDataPoint(AgentDataPoint):
    scorer_criteria: ScorerCriteria


class ScorerFacingDataset(BaseModel):
    dataset: list[ScorerDataPoint]
