from datetime import datetime
from typing import Any, Literal, Dict

from pydantic import BaseModel, Field

class ModelMetrics(BaseModel):
    # Identification
    session_id: str = Field(description="Session ID")
    query_id: str = Field(description="Query ID")

    # Model info
    model: str = Field(description="Model name", default=None)
    model_provider: str | None = Field(description="Model provider", default=None)
    created_at: datetime | None = Field(default=None, description="Model timestamp")

    total_duration: int | None = Field(default=None, description="Total duration (ns)")
    load_duration: int | None = Field(default=None, description="Load duration (ns)")
    prompt_eval_duration: int | None = Field(default=None, description="Prompt eval duration (ns)")
    eval_duration: int | None = Field(default=None, description="Eval duration (ns)")

    prompt_eval_count: int | None = Field(default=None, description="Prompt eval tokens")
    eval_count: int | None = Field(default=None, description="Eval tokens")
    input_tokens: int | None = Field(default=None, description="Input tokens")
    output_tokens: int | None = Field(default=None, description="Output tokens")
    total_tokens: int | None = Field(default=None, description="Total tokens")

    # Raw metadata
    raw_metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Unmodified metadata from model + LangGraph"
    )


class ServerConfig(BaseModel):
    """Server Configuration."""

    host: str
    port: int
    transport: Literal["stdio", "sse", "streamable-http"]
    url: str


class PlannerTask(BaseModel):
    """Represents a single task generated by the Planner."""

    id: int = Field(description="Sequential ID for the task.")
    description: str = Field(
        description="Clear description of the task to be executed."
    )
    status: (
        Any
        | Literal[
            "input_required",
            "completed",
            "error",
            "pending",
            "incomplete",
            "todo",
            "not_started",
        ]
        | None
    ) = Field(description="Status of the task", default="input_required")


class ModelInfo(BaseModel):
    """OpenStudio model Info."""

    model_path: str | None = Field(
        description="=The local directory of the building energy model file."
    )
    building_type: Literal["large-office", "medium-office", "small-office"] = (
        "small-office"
    )
    window_to_wall_ratio: str | None = Field(
        description="Total window to wall ratio of the model"
    )


class TaskList(BaseModel):
    """Output schema for the Planner Agent"""

    original_query: str | None = Field(
        description="The original user query for context."
    )
    blackboard: ModelInfo | None = Field(description="Task information")

    tasks: list[PlannerTask] = Field(
        description="A list of tasks to be executed sequentially."
    )
