models.py

"""Data models for Trinka Agent Backend"""
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Literal


# Request Models
class TextAnalysisRequest(BaseModel):
    """Request model for text analysis"""
    text: str = Field(..., description="Text content to analyze")
    mode: Literal["power", "lite"] = Field(default="power", description="Analysis mode")
    language: str = Field(default="en-US", description="Language code")


class GrammarCheckRequest(BaseModel):
    """Request model for grammar checking"""
    text: str = Field(..., description="Text to check for grammar errors")


class ParaphraseRequest(BaseModel):
    """Request model for paraphrasing"""
    text: str = Field(..., description="Text to paraphrase")
    style: Literal["academic", "concise", "simplified"] = Field(
        default="academic", 
        description="Paraphrase style"
    )


class ConsistencyCheckRequest(BaseModel):
    """Request model for consistency checking"""
    text: str = Field(..., description="Document text to check for consistency")


# Response Models
class Position(BaseModel):
    """Text position for highlighting"""
    start: int = Field(..., description="Start character position")
    end: int = Field(..., description="End character position")


class GrammarIssue(BaseModel):
    """Individual grammar issue"""
    id: str = Field(..., description="Unique identifier")
    error_type: str = Field(..., description="Type of error (capitalization, punctuation, etc.)")
    original_text: str = Field(..., description="Original text with error")
    suggested_fix: str = Field(..., description="Suggested correction")
    message: str = Field(..., description="Explanation of the issue")
    position: Position = Field(..., description="Position in text")


class GrammarCheckResponse(BaseModel):
    """Response model for grammar checking"""
    tool: str = Field(default="grammar_checker")
    issues: List[GrammarIssue] = Field(default_factory=list)
    total_issues: int = Field(default=0)


class ParaphraseResponse(BaseModel):
    """Response model for paraphrasing"""
    tool: str = Field(default="paraphraser")
    original_text: str = Field(..., description="Original text")
    paraphrased_text: str = Field(..., description="Rewritten text")
    style: str = Field(..., description="Style used for paraphrasing")


class ConsistencyIssue(BaseModel):
    """Individual consistency issue"""
    id: str = Field(..., description="Unique identifier")
    category: str = Field(..., description="Category (spelling, style, dashes, symbols)")
    term: str = Field(..., description="Inconsistent term")
    variant_found: str = Field(..., description="Variant found in text")
    suggested_standard: str = Field(..., description="Suggested standard format")
    occurrences: int = Field(default=1, description="Number of occurrences")


class ConsistencyCheckResponse(BaseModel):
    """Response model for consistency checking"""
    tool: str = Field(default="consistency_checker")
    issues: List[ConsistencyIssue] = Field(default_factory=list)
    total_issues: int = Field(default=0)


# Report Models
class CitationRiskGroup(BaseModel):
    """Citation risk group"""
    risk_level: str = Field(..., description="Risk level (high, medium, low, none)")
    count: int = Field(..., description="Number of citations in this category")
    description: str = Field(..., description="Description of the risk")


class CitationReportResponse(BaseModel):
    """Response model for citation report"""
    tool: str = Field(default="citation_checker")
    total_citations: int = Field(..., description="Total number of citations")
    risk_groups: List[CitationRiskGroup] = Field(default_factory=list)
    summary: str = Field(..., description="Summary of citation analysis")


class PlagiarismSource(BaseModel):
    """Plagiarism source"""
    url: str = Field(..., description="Source URL")
    match_percentage: str = Field(..., description="Match percentage")


class PlagiarismReportResponse(BaseModel):
    """Response model for plagiarism report"""
    tool: str = Field(default="plagiarism_checker")
    overall_score: str = Field(..., description="Overall similarity score")
    internet_match: str = Field(..., description="Internet match percentage")
    publication_match: str = Field(..., description="Publication match percentage")
    top_sources: List[PlagiarismSource] = Field(default_factory=list)


class JournalSuggestion(BaseModel):
    """Journal suggestion"""
    name: str = Field(..., description="Journal name")
    confidence: str = Field(..., description="Confidence score")
    impact_factor: Optional[str] = Field(None, description="Impact factor")


class JournalFinderResponse(BaseModel):
    """Response model for journal finder"""
    tool: str = Field(default="journal_finder")
    suggestions: List[JournalSuggestion] = Field(default_factory=list)
    total_suggestions: int = Field(default=0)


# Agent Models
class Alert(BaseModel):
    """Alert model for UI integration"""
    id: str = Field(..., description="Unique identifier")
    type: str = Field(..., description="Alert type (grammar, style, spelling, etc.)")
    title: str = Field(..., description="Alert title")
    beforeText: str = Field(..., description="Original text")
    afterText: str = Field(..., description="Suggested text")
    message: str = Field(..., description="Explanation")
    similarCount: Optional[int] = Field(None, description="Number of similar alerts")
    position: Position = Field(..., description="Position in text")


class AgentResponse(BaseModel):
    """Response model for agent analysis"""
    intent: str = Field(..., description="Detected intent")
    tool: str = Field(..., description="Tool used")
    alerts: List[Alert] = Field(default_factory=list)
    total_alerts: int = Field(default=0)
    processing_time: float = Field(..., description="Processing time in seconds")
    confidence: float = Field(default=1.0, description="Confidence score (0-1)")


# Chat Agent Models (as per PRD)
class AgentActionButton(BaseModel):
    """Action button in agent response"""
    id: str = Field(..., description="Button identifier")
    label: str = Field(..., description="Button label text")
    action: str = Field(..., description="Action to perform")
    variant: Literal["primary", "secondary", "tertiary"] = Field(
        default="secondary", description="Button style variant"
    )


class AlertGroupedItem(BaseModel):
    """Individual alert within a group"""
    id: str = Field(..., description="Alert identifier")
    beforeText: str = Field(..., description="Original text")
    afterText: str = Field(..., description="Suggested text")
    context: Optional[str] = Field(None, description="Context or explanation")
    position: Optional[Position] = Field(None, description="Position in text")


class AlertGroup(BaseModel):
    """Categorized group of similar alerts"""
    id: str = Field(..., description="Group identifier")
    category: str = Field(..., description="Category name (e.g., 'Fix all typos')")
    categoryType: str = Field(..., description="Category type")
    toolCategory: str = Field(..., description="Tool category (grammar, paraphraser, etc.)")
    count: int = Field(..., description="Number of alerts in group")
    isExpanded: bool = Field(default=False, description="Whether group is expanded in UI")
    alerts: List[AlertGroupedItem] = Field(default_factory=list, description="Alerts in this group")


class ChatContext(BaseModel):
    """Context for chat requests"""
    documentText: str = Field(..., description="Full document text")
    conversationHistory: List[Dict[str, Any]] = Field(
        default_factory=list, description="Previous messages in conversation"
    )
    activeTab: Optional[str] = Field(None, description="Currently active tab")
    conversationId: Optional[str] = Field(None, description="Conversation ID for persistence")


class AgentChatRequest(BaseModel):
    """Request model for /api/agent/chat endpoint"""
    message: str = Field(..., description="User message")
    context: ChatContext = Field(..., description="Conversation context")


class AgentChatResponse(BaseModel):
    """Response model for /api/agent/chat endpoint"""
    message: str = Field(..., description="Agent's response message")
    buttons: List[AgentActionButton] = Field(
        default_factory=list, description="Action buttons to display"
    )
    alertGroups: List[AlertGroup] = Field(
        default_factory=list, description="Categorized alert groups"
    )
    conversationId: Optional[str] = Field(None, description="Conversation ID")
    timestamp: Optional[str] = Field(None, description="Response timestamp")


class AppliedChange(BaseModel):
    """Individual change that was applied"""
    original: str = Field(..., description="Original text")
    corrected: str = Field(..., description="Corrected text")
    position: Position = Field(..., description="Position in text")


class ActionContext(BaseModel):
    """Context for action requests"""
    documentText: str = Field(..., description="Current document text")
    alertGroupId: Optional[str] = Field(None, description="Alert group ID if applicable")
    conversationId: Optional[str] = Field(None, description="Conversation ID")


class AgentActionRequest(BaseModel):
    """Request model for /api/agent/action endpoint"""
    action: str = Field(..., description="Action to perform")
    context: ActionContext = Field(..., description="Action context")


class AgentActionResponse(BaseModel):
    """Response model for /api/agent/action endpoint"""
    updatedText: str = Field(..., description="Updated document text")
    appliedChanges: List[AppliedChange] = Field(
        default_factory=list, description="List of applied changes"
    )
    nextMessage: Optional[str] = Field(None, description="Follow-up message from agent")
    conversationId: Optional[str] = Field(None, description="Conversation ID")