from typing import Any

from pydantic import BaseModel, Field

from app.core.constants import ActorValue
from app.modules.ai_agent.constants import AIRiskPreference
from app.modules.audit.constants import AuditSource


class AIAskRequest(BaseModel):
    prompt: str
    context: dict[str, Any] = Field(default_factory=dict)
    actor: str = ActorValue.API
    source: str = AuditSource.API


class AIAskResponse(BaseModel):
    ok: bool = True
    provider: str
    answer: str
    raw: dict[str, Any] = Field(default_factory=dict)
    error: str | None = None


class DraftPolicyRequest(BaseModel):
    title: str
    policy_type: str
    requirements: list[str]
    actor: str = ActorValue.API


class InvestmentResearchRequest(BaseModel):
    symbol_or_topic: str
    risk_preference: str = AIRiskPreference.BALANCED
    actor: str = ActorValue.API