app/copilot/schemas.py

from __future__ import annotations

from datetime import datetime
from typing import Annotated, Literal
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field, model_validator

from app.social.schemas.posts import SocialPostCreate
from app.social.schemas.scheduling import SocialScheduleCreate
from app.templates.marketplace_schemas import SlotBinding

CopilotRunStatus = Literal[
    "plan_ready",
    "blocked",
    "executing",
    "completed",
    "partial",
    "failed",
    "cancelled",
]
CopilotActionStatus = Literal["pending", "running", "completed", "failed", "cancelled"]


class CopilotEditorSummary(BaseModel):
    model_config = ConfigDict(extra="forbid")

    revision: int = Field(ge=1)
    duration_ms: int = Field(ge=0)
    track_count: int = Field(ge=0, le=32)
    clip_count: int = Field(ge=0, le=500)


class CopilotContextInput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    project_id: UUID | None = None
    selected_asset_ids: list[UUID] = Field(default_factory=list, max_length=20)
    selected_clip_ids: list[str] = Field(default_factory=list, max_length=20)
    active_tool: str | None = Field(default=None, max_length=100)
    editor_summary: CopilotEditorSummary | None = None


class CopilotContext(CopilotContextInput):
    workspace_id: str
    available_capabilities: list[str] = Field(default_factory=list, max_length=100)


class ProjectOpenArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID


class AssetSelectArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    asset_id: UUID
    project_id: UUID | None = None


class AiGenerateImageArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    prompt: str = Field(min_length=1, max_length=4_000)
    project_id: UUID | None = None
    source_asset_id: UUID | None = None
    model: str | None = Field(default=None, max_length=255)


class AiGenerateVideoArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    prompt: str = Field(min_length=1, max_length=4_000)
    project_id: UUID | None = None
    source_asset_id: UUID
    model: str | None = Field(default=None, max_length=255)


class EditorSplitClipArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID
    clip_id: str = Field(min_length=1, max_length=128)
    at_ms: int = Field(gt=0)
    expected_revision: int = Field(ge=1)


class EditorDeleteClipArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID
    clip_id: str = Field(min_length=1, max_length=128)
    expected_revision: int = Field(ge=1)


class EditorSetDurationArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID
    clip_id: str = Field(min_length=1, max_length=128)
    duration_ms: int = Field(gt=0)
    expected_revision: int = Field(ge=1)


class EditorAddClipArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID
    asset_id: UUID
    expected_revision: int = Field(ge=1)
    duration_ms: int = Field(default=5_000, gt=0)


class EditorRenderArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID
    expected_revision: int = Field(ge=1)

class TemplateSearchArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    query: str = Field(min_length=1, max_length=200)
    category: str | None = Field(default=None, max_length=50)


class TemplateGetArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    template_id: UUID


class TemplateApplyArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    template_id: UUID
    project_id: UUID
    template_version_id: UUID | None = None
    slot_bindings: dict[str, SlotBinding] = Field(default_factory=dict, max_length=100)


class TemplateCreateProjectArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    template_id: UUID
    project_name: str = Field(min_length=1, max_length=200)
    template_version_id: UUID | None = None
    slot_bindings: dict[str, SlotBinding] = Field(default_factory=dict, max_length=100)


class PublishingValidateArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    post_id: UUID


class PublishingCreatePostArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    post: SocialPostCreate


class PublishingScheduleArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    post_id: UUID
    schedule: SocialScheduleCreate


class PublishingPublishArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    post_id: UUID


class PublishingCancelArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    post_id: UUID


class AnalyticsOverviewArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID | None = None
    provider: str | None = Field(default=None, max_length=32)
    metric: Literal["views", "impressions", "likes", "comments", "shares", "engagement_rate"] = (
        "views"
    )
    timezone: str = Field(default="UTC", min_length=1, max_length=100)


class AnalyticsSyncArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: UUID | None = None
    provider: str | None = Field(default=None, max_length=32)
    timezone: str = Field(default="UTC", min_length=1, max_length=100)


class _ActionBase(BaseModel):
    model_config = ConfigDict(extra="forbid")

    id: str = Field(min_length=1, max_length=64)
    reason: str = Field(min_length=1, max_length=500)
    requires_confirmation: bool
    destructive: bool
    external_side_effect: bool
    required_permission: str = Field(min_length=1, max_length=100)
    required_capability: str = Field(min_length=1, max_length=100)


class ProjectOpenAction(_ActionBase):
    type: Literal["project.open"]
    arguments: ProjectOpenArguments


class AssetSelectAction(_ActionBase):
    type: Literal["asset.select"]
    arguments: AssetSelectArguments


class AiGenerateImageAction(_ActionBase):
    type: Literal["ai.generate_image"]
    arguments: AiGenerateImageArguments


class AiGenerateVideoAction(_ActionBase):
    type: Literal["ai.generate_video"]
    arguments: AiGenerateVideoArguments


class EditorSplitClipAction(_ActionBase):
    type: Literal["editor.split_clip"]
    arguments: EditorSplitClipArguments


class EditorDeleteClipAction(_ActionBase):
    type: Literal["editor.delete_clip"]
    arguments: EditorDeleteClipArguments


class EditorSetDurationAction(_ActionBase):
    type: Literal["editor.set_duration"]
    arguments: EditorSetDurationArguments


class EditorAddClipAction(_ActionBase):
    type: Literal["editor.add_clip"]
    arguments: EditorAddClipArguments


class EditorRenderAction(_ActionBase):
    type: Literal["editor.render"]
    arguments: EditorRenderArguments


class ShareProjectArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: str
    user_id: str
    role: str

class SubmitForReviewArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: str
    workflow_id: str

class ProjectActivitySummarizeArguments(BaseModel):
    model_config = ConfigDict(extra="forbid")
    project_id: str

class ShareProjectAction(_ActionBase):
    type: Literal["project.share"]
    arguments: ShareProjectArguments

class SubmitForReviewAction(_ActionBase):
    type: Literal["review.submit"]
    arguments: SubmitForReviewArguments

class ProjectActivitySummarizeAction(_ActionBase):
    type: Literal["project_activity.summarize"]
    arguments: ProjectActivitySummarizeArguments


class TemplateSearchAction(_ActionBase):
    type: Literal["template.search"]
    arguments: TemplateSearchArguments


class TemplateGetAction(_ActionBase):
    type: Literal["template.get"]
    arguments: TemplateGetArguments


class TemplateApplyAction(_ActionBase):
    type: Literal["template.apply"]
    arguments: TemplateApplyArguments


class TemplateCreateProjectAction(_ActionBase):
    type: Literal["template.create_project"]
    arguments: TemplateCreateProjectArguments


class PublishingValidateAction(_ActionBase):
    type: Literal["publishing.validate"]
    arguments: PublishingValidateArguments


class PublishingCreatePostAction(_ActionBase):
    type: Literal["publishing.create_post"]
    arguments: PublishingCreatePostArguments


class PublishingScheduleAction(_ActionBase):
    type: Literal["publishing.schedule"]
    arguments: PublishingScheduleArguments


class PublishingPublishAction(_ActionBase):
    type: Literal["publishing.publish"]
    arguments: PublishingPublishArguments


class PublishingCancelAction(_ActionBase):
    type: Literal["publishing.cancel"]
    arguments: PublishingCancelArguments


class AnalyticsOverviewAction(_ActionBase):
    type: Literal["analytics.overview"]
    arguments: AnalyticsOverviewArguments


class AnalyticsSyncAction(_ActionBase):
    type: Literal["analytics.sync"]
    arguments: AnalyticsSyncArguments


CopilotAction = Annotated[
    ProjectOpenAction
    | AssetSelectAction
    | AiGenerateImageAction
    | AiGenerateVideoAction
    | EditorSplitClipAction
    | EditorDeleteClipAction
    | EditorSetDurationAction
    | EditorAddClipAction
    | EditorRenderAction
    | ShareProjectAction
    | SubmitForReviewAction
    | ProjectActivitySummarizeAction
    | TemplateSearchAction
    | TemplateGetAction
    | TemplateApplyAction
    | TemplateCreateProjectAction
    | PublishingValidateAction
    | PublishingCreatePostAction
    | PublishingScheduleAction
    | PublishingPublishAction
    | PublishingCancelAction
    | AnalyticsOverviewAction
    | AnalyticsSyncAction,
    Field(discriminator="type"),
]


class CopilotPlan(BaseModel):
    model_config = ConfigDict(extra="forbid")

    intent: str = Field(min_length=1, max_length=200)
    explanation: str = Field(min_length=1, max_length=1_000)
    actions: list[CopilotAction] = Field(default_factory=list, max_length=20)
    missing_information: list[str] = Field(default_factory=list, max_length=20)
    unsupported_capabilities: list[str] = Field(default_factory=list, max_length=20)
    executable: bool
    requires_confirmation: bool

    @model_validator(mode="after")
    def validate_execution(self) -> "CopilotPlan":
        if self.executable and not self.actions:
            raise ValueError("Executable plans require at least one action")
        if self.requires_confirmation != any(
            action.requires_confirmation for action in self.actions
        ):
            raise ValueError("Plan confirmation state must match its actions")
        return self


class CopilotActionResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    action_id: str
    action_type: str
    status: CopilotActionStatus
    summary: str = Field(max_length=1_000)
    resource_type: str | None = Field(default=None, max_length=100)
    resource_id: str | None = Field(default=None, max_length=255)
    retryable: bool = False
    error_code: str | None = Field(default=None, max_length=100)


class CopilotRunCreate(BaseModel):
    model_config = ConfigDict(extra="forbid", str_strip_whitespace=True)

    request: str = Field(min_length=1, max_length=4_000)
    context: CopilotContextInput = Field(default_factory=CopilotContextInput)


class CopilotExecuteRequest(BaseModel):
    model_config = ConfigDict(extra="forbid")
    confirmed: bool = False


class CopilotRun(BaseModel):
    model_config = ConfigDict(extra="forbid")

    id: str
    project_id: str | None
    status: CopilotRunStatus
    request: str
    context: CopilotContext
    plan: CopilotPlan
    current_action_id: str | None
    results: list[CopilotActionResult]
    summary: str | None
    error_code: str | None
    error_message: str | None
    created_at: datetime
    updated_at: datetime
    completed_at: datetime | None


class CopilotRunList(BaseModel):
    items: list[CopilotRun]
    offset: int
    limit: int


class CopilotActionCapability(BaseModel):
    model_config = ConfigDict(extra="forbid")
    type: str
    description: str
    required_permission: str
    required_capability: str
    destructive: bool
    external_side_effect: bool
    requires_confirmation: bool
    available: bool


class CopilotCapabilities(BaseModel):
    model_config = ConfigDict(extra="forbid")
    available: bool
    planner: Literal["deterministic"]
    actions: list[CopilotActionCapability]
    permissions: list[str]