"""Analysis and research output models."""

from pydantic import BaseModel, Field


class DataPoint(BaseModel):
    """A single data point for analysis."""

    label: str
    value: float
    metadata: dict[str, str] = {}


class AnalysisResult(BaseModel):
    """Output from the analyst agent."""

    summary: str = Field(description="Brief summary of the analysis")
    confidence: float = Field(ge=0, le=1, description="Confidence score 0-1")
    data_points: list[DataPoint] = Field(default_factory=list)
    recommendations: list[str] = Field(default_factory=list)
    methodology: str | None = Field(default=None, description="How the analysis was performed")


class ResearchResult(BaseModel):
    """Output from the research agent."""

    answer: str = Field(description="The researched answer")
    sources: list[str] = Field(description="URLs or references used")
    confidence: float = Field(ge=0, le=1, default=0.8)
    follow_up_questions: list[str] = Field(
        default_factory=list, description="Suggested follow-up questions"
    )

    def format_with_sources(self) -> str:
        """Format the answer with source citations."""
        sources_text = "\n".join(f"- {s}" for s in self.sources)
        return f"{self.answer}\n\nSources:\n{sources_text}"
