from pydantic import BaseModel, Field
from typing import List, Optional
from datetime import datetime
from bson import ObjectId

class QuestionConfig(BaseModel):
    type: str = Field(..., description="mcq, short_answer, long_answer, true_false")
    count: int = Field(..., ge=1)
    marksPerQuestion: int = Field(..., ge=1)
    difficulty: str = Field(..., description="easy, medium, hard")

class AssignmentInput(BaseModel):
    title: str = Field(..., max_length=200)
    subject: str = Field(..., max_length=100)
    grade: str = Field(..., max_length=20)
    schoolName: str = Field(..., max_length=200)
    dueDate: str
    totalMarks: int = Field(..., ge=1)
    duration: int = Field(..., ge=1)
    questionConfigs: List[QuestionConfig] = Field(..., min_length=1)
    additionalInstructions: Optional[str] = Field(None, max_length=2000)
    uploadedFileContent: Optional[str] = None

class QuestionSubSchema(BaseModel):
    id: str
    text: str
    type: str
    difficulty: str
    marks: int
    options: Optional[List[str]] = None
    answerKey: str

class SectionSubSchema(BaseModel):
    id: str
    title: str
    instruction: str
    questionType: str
    questions: List[QuestionSubSchema]
    totalMarks: int

class GeneratedPaperSubSchema(BaseModel):
    id: str
    assignmentId: str
    schoolName: str
    subject: str
    grade: str
    duration: int
    totalMarks: int
    sections: List[SectionSubSchema]
    createdAt: str

class AssignmentDocument(BaseModel):
    id: Optional[str] = Field(None, alias="_id")
    input: AssignmentInput
    status: str = Field(default="pending", description="pending, processing, completed, failed")
    jobId: Optional[str] = None
    result: Optional[GeneratedPaperSubSchema] = None
    error: Optional[str] = None
    createdAt: Optional[datetime] = None
    updatedAt: Optional[datetime] = None

    class Config:
        populate_by_name = True
        json_encoders = {ObjectId: str}