from pydantic import BaseModel, Field
from typing import Optional, List, Literal, Union, Dict, Any, Callable
from enum import Enum


class DeepResearchMode(str, Enum):
    """Research mode options."""

    FAST = "fast"
    STANDARD = "standard"
    LITE = "lite"  # Deprecated: use STANDARD instead (kept for backward compatibility)
    HEAVY = "heavy"


class DeepResearchStatus(str, Enum):
    """Task status options."""

    QUEUED = "queued"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class FileAttachment(BaseModel):
    """File attachment for research."""

    data: str = Field(..., description="Data URL (base64 encoded)")
    filename: str = Field(..., description="Original filename")
    media_type: str = Field(..., description="MIME type")
    context: Optional[str] = Field(None, description="Context about the file")


class MCPServerConfig(BaseModel):
    """MCP server configuration."""

    url: str = Field(..., description="MCP server URL")
    name: Optional[str] = Field(None, description="Server name")
    tool_prefix: Optional[str] = Field(None, description="Custom tool prefix")
    auth: Optional[Dict[str, Any]] = Field(None, description="Authentication config")
    allowed_tools: Optional[List[str]] = Field(None, description="Allowed tools")


class Deliverable(BaseModel):
    """Deliverable file configuration."""

    type: Literal["csv", "xlsx", "pptx", "docx", "pdf"] = Field(
        ..., description="File type"
    )
    description: str = Field(
        ..., max_length=500, description="What data to extract or content to generate"
    )
    columns: Optional[List[str]] = Field(
        None, description="Suggested column names (for CSV/XLSX)"
    )
    include_headers: Optional[bool] = Field(
        True, description="Include column headers (for CSV/XLSX)"
    )
    sheet_name: Optional[str] = Field(None, description="Sheet name (for XLSX only)")
    slides: Optional[int] = Field(None, description="Number of slides (for PPTX only)")
    template: Optional[str] = Field(None, description="Template name to use")


class DeliverableResult(BaseModel):
    """Deliverable generation result."""

    id: str = Field(..., description="Unique deliverable ID")
    request: str = Field(..., description="Original request description")
    type: Literal["csv", "xlsx", "pptx", "docx", "pdf", "unknown"] = Field(
        ..., description="Deliverable file type"
    )
    status: Literal["completed", "failed"] = Field(..., description="Generation status")
    title: str = Field(..., description="Generated filename/title")
    description: Optional[str] = Field(
        None, description="Deliverable content description"
    )
    url: str = Field(
        ..., description="Token-signed authenticated URL to download the file"
    )
    s3_key: str = Field(..., description="S3 storage key")
    row_count: Optional[int] = Field(None, description="Number of rows (for CSV/XLSX)")
    column_count: Optional[int] = Field(
        None, description="Number of columns (for CSV/XLSX)"
    )
    error: Optional[str] = Field(None, description="Error message if status is failed")
    created_at: int = Field(..., description="Unix timestamp of creation")


class SearchConfig(BaseModel):
    """Search configuration for Deep Research tasks and batches.

    Controls which data sources are queried, what content is included/excluded,
    and how results are filtered by date or category.

    Attributes:
        search_type: Controls which backend search systems are queried.
            - "all" (default): Searches both web and proprietary data sources
            - "web": Searches only web sources (general web search, news, articles)
            - "proprietary": Searches only proprietary/internal data sources
        included_sources: Restricts search to only the specified source types.
            Available types: "web", "academic", "finance", "patent",
            "transportation", "politics", "legal"
        excluded_sources: Excludes specific source types from search results.
            Uses the same source type values as included_sources.
            Cannot be used simultaneously with included_sources.
        start_date: Filters results to content published on or after this date.
            Format: ISO date (YYYY-MM-DD), e.g., "2024-01-01"
        end_date: Filters results to content published on or before this date.
            Format: ISO date (YYYY-MM-DD), e.g., "2024-12-31"
        category: Filters results by a specific category.
            Category values are source-dependent.
    """

    search_type: Optional[Literal["all", "web", "proprietary"]] = Field(
        None, description="Search scope: 'all', 'web', or 'proprietary'"
    )
    included_sources: Optional[List[str]] = Field(
        None,
        description="Source types to include: 'web', 'academic', 'finance', 'patent', 'transportation', 'politics', 'legal'",
    )
    excluded_sources: Optional[List[str]] = Field(
        None, description="Source types to exclude from search"
    )
    start_date: Optional[str] = Field(
        None, description="Start date filter in ISO format (YYYY-MM-DD)"
    )
    end_date: Optional[str] = Field(
        None, description="End date filter in ISO format (YYYY-MM-DD)"
    )
    category: Optional[str] = Field(None, description="Category filter for results")


class Progress(BaseModel):
    """Task progress information."""

    current_step: int
    total_steps: int


class ChartDataPoint(BaseModel):
    """Chart data point."""

    x: Union[str, int, float]
    y: Union[int, float]


class ChartDataSeries(BaseModel):
    """Chart data series."""

    name: str
    data: List[ChartDataPoint]


class ImageMetadata(BaseModel):
    """Image metadata."""

    image_id: str
    image_type: Literal["chart", "ai_generated"]
    deepresearch_id: str
    title: str
    description: Optional[str] = None
    image_url: str
    s3_key: str
    created_at: int
    chart_type: Optional[Literal["line", "bar", "area"]] = None
    x_axis_label: Optional[str] = None
    y_axis_label: Optional[str] = None
    data_series: Optional[List[ChartDataSeries]] = None


class DeepResearchSource(BaseModel):
    """Source information."""

    title: str
    url: str
    snippet: Optional[str] = None
    description: Optional[str] = None
    source: Optional[str] = None
    org_id: Optional[str] = None
    price: Optional[float] = None
    id: Optional[str] = None
    doc_id: Optional[int] = None
    doi: Optional[str] = None
    category: Optional[str] = None
    source_id: Optional[int] = None
    word_count: Optional[int] = None


class Usage(BaseModel):
    """Usage and cost information."""

    search_cost: float
    contents_cost: float
    ai_cost: float
    compute_cost: float
    total_cost: float


class DeepResearchCreateResponse(BaseModel):
    """Response from creating a deep research task."""

    success: bool
    deepresearch_id: Optional[str] = None
    status: Optional[DeepResearchStatus] = None
    model: Optional[DeepResearchMode] = None
    created_at: Optional[str] = None
    metadata: Optional[Dict[str, Union[str, int, bool]]] = None
    public: Optional[bool] = None
    webhook_secret: Optional[str] = None
    message: Optional[str] = None
    error: Optional[str] = None


class DeepResearchStatusResponse(BaseModel):
    """Response from getting task status."""

    success: bool
    deepresearch_id: Optional[str] = None
    status: Optional[DeepResearchStatus] = None
    query: Optional[str] = None
    mode: Optional[DeepResearchMode] = None
    output_formats: Optional[
        List[Union[Literal["markdown", "pdf"], Dict[str, Any]]]
    ] = None
    created_at: Optional[int] = None
    public: Optional[bool] = None

    # Optional fields based on status
    progress: Optional[Progress] = None
    messages: Optional[List[Any]] = None
    completed_at: Optional[int] = None
    output: Optional[Union[str, Dict[str, Any], Any]] = None
    output_type: Optional[Literal["markdown", "json"]] = None
    pdf_url: Optional[str] = None
    images: Optional[List[ImageMetadata]] = None
    deliverables: Optional[List[DeliverableResult]] = None
    sources: Optional[List[DeepResearchSource]] = None
    cost: Optional[float] = None
    error: Optional[str] = None


class DeepResearchTaskListItem(BaseModel):
    """Minimal task info for list view."""

    deepresearch_id: str
    query: str
    status: DeepResearchStatus
    created_at: int
    public: Optional[bool] = None


class DeepResearchListResponse(BaseModel):
    """Response from listing tasks."""

    success: bool
    data: Optional[List[Dict[str, Any]]] = None
    error: Optional[str] = None


class DeepResearchUpdateResponse(BaseModel):
    """Response from updating a task."""

    success: bool
    message: Optional[str] = None
    deepresearch_id: Optional[str] = None
    error: Optional[str] = None


class DeepResearchCancelResponse(BaseModel):
    """Response from cancelling a task."""

    success: bool
    message: Optional[str] = None
    deepresearch_id: Optional[str] = None
    error: Optional[str] = None


class DeepResearchDeleteResponse(BaseModel):
    """Response from deleting a task."""

    success: bool
    message: Optional[str] = None
    deepresearch_id: Optional[str] = None
    error: Optional[str] = None


class DeepResearchTogglePublicResponse(BaseModel):
    """Response from toggling public flag."""

    success: bool
    message: Optional[str] = None
    deepresearch_id: Optional[str] = None
    public: Optional[bool] = None
    error: Optional[str] = None


# =============================================================================
# Batch Types
# =============================================================================


class BatchStatus(str, Enum):
    """Batch status options."""

    OPEN = "open"
    PROCESSING = "processing"
    COMPLETED = "completed"
    COMPLETED_WITH_ERRORS = "completed_with_errors"
    CANCELLED = "cancelled"


class BatchCounts(BaseModel):
    """Task counts within a batch."""

    total: int = Field(..., description="Total number of tasks")
    queued: int = Field(..., description="Number of queued tasks")
    running: int = Field(..., description="Number of running tasks")
    completed: int = Field(..., description="Number of completed tasks")
    failed: int = Field(..., description="Number of failed tasks")
    cancelled: int = Field(..., description="Number of cancelled tasks")


class BatchUsage(BaseModel):
    """Aggregated usage for a batch."""

    search_cost: float = Field(..., description="Total search cost")
    contents_cost: float = Field(..., description="Total contents cost")
    ai_cost: float = Field(..., description="Total AI cost")
    total_cost: float = Field(..., description="Total cost")


class BatchTaskInput(BaseModel):
    """Input for a batch task."""

    id: Optional[str] = Field(None, description="User-provided task ID")
    input: str = Field(..., description="Research query or task description")
    strategy: Optional[str] = Field(None, description="Natural language strategy")
    urls: Optional[List[str]] = Field(None, description="URLs to extract and analyze")
    metadata: Optional[Dict[str, Union[str, int, bool]]] = Field(
        None, description="Custom metadata"
    )


class DeepResearchBatch(BaseModel):
    """Batch of deep research tasks."""

    batch_id: str = Field(..., description="Unique batch ID")
    organisation_id: Optional[str] = Field(None, description="Organization ID")
    api_key_id: Optional[str] = Field(None, description="API key ID")
    credit_id: Optional[str] = Field(None, description="Credit ID")
    name: Optional[str] = Field(None, description="Batch name")
    status: BatchStatus = Field(..., description="Current batch status")
    model: DeepResearchMode = Field(..., description="Default model for tasks")
    output_formats: Optional[
        List[Union[Literal["markdown", "pdf"], Dict[str, Any]]]
    ] = Field(None, description="Default output formats")
    search_params: Optional[Dict[str, Any]] = Field(
        None, description="Default search parameters"
    )
    counts: BatchCounts = Field(..., description="Task counts")
    usage: Optional[BatchUsage] = Field(None, description="Aggregated usage")
    webhook_url: Optional[str] = Field(
        None, description="Webhook URL for notifications"
    )
    webhook_secret: Optional[str] = Field(None, description="Webhook secret")
    created_at: Union[int, str] = Field(..., description="Creation timestamp")
    updated_at: Optional[Union[int, str]] = Field(
        None, description="Last update timestamp"
    )
    completed_at: Optional[Union[int, str]] = Field(
        None, description="Completion timestamp"
    )
    metadata: Optional[Dict[str, Union[str, int, bool]]] = Field(
        None, description="Custom metadata"
    )


class BatchCreateResponse(BaseModel):
    """Response from creating a batch."""

    success: bool
    batch_id: Optional[str] = None
    status: Optional[BatchStatus] = None
    model: Optional[DeepResearchMode] = None
    counts: Optional[BatchCounts] = None
    created_at: Optional[Union[int, str]] = None
    webhook_secret: Optional[str] = None
    message: Optional[str] = None
    error: Optional[str] = None


class BatchAddTasksResponse(BaseModel):
    """Response from adding tasks to a batch."""

    success: bool
    batch_id: Optional[str] = None
    added: Optional[int] = None
    task_ids: Optional[List[str]] = None
    message: Optional[str] = None
    error: Optional[str] = None


class BatchStatusResponse(BaseModel):
    """Response from getting batch status."""

    success: bool
    batch: Optional[DeepResearchBatch] = None
    error: Optional[str] = None


class BatchTaskListItem(BaseModel):
    """Minimal task info in batch list."""

    deepresearch_id: str
    batch_task_id: Optional[str] = None
    query: str
    status: DeepResearchStatus
    created_at: Union[int, str]
    completed_at: Optional[Union[int, str]] = None
    error: Optional[str] = None


class BatchTasksListResponse(BaseModel):
    """Response from listing tasks in a batch."""

    success: bool
    batch_id: Optional[str] = None
    tasks: Optional[List[BatchTaskListItem]] = None
    error: Optional[str] = None


class BatchCancelResponse(BaseModel):
    """Response from cancelling a batch."""

    success: bool
    batch_id: Optional[str] = None
    message: Optional[str] = None
    error: Optional[str] = None


class BatchListResponse(BaseModel):
    """Response from listing batches."""

    success: bool
    batches: Optional[List[DeepResearchBatch]] = None
    error: Optional[str] = None
