import enum
import uuid

from sqlalchemy import Enum, ForeignKey, Integer, String, Text
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import Mapped, mapped_column, relationship

from app.db.session import Base
from app.models.base import TimestampMixin, UUIDPKMixin


class ExampleType(str, enum.Enum):
    qa = "qa"
    conversation = "conversation"
    scenario = "scenario"
    reasoning = "reasoning"
    diagnostic = "diagnostic"
    follow_up = "follow_up"


class Difficulty(str, enum.Enum):
    beginner = "beginner"
    intermediate = "intermediate"
    advanced = "advanced"


class ExampleStatus(str, enum.Enum):
    pending = "pending"
    approved = "approved"
    rejected = "rejected"
    edited = "edited"


class Dataset(Base, UUIDPKMixin, TimestampMixin):
    __tablename__ = "datasets"

    project_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("projects.id"), nullable=False)
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[str | None] = mapped_column(Text, nullable=True)
    version: Mapped[int] = mapped_column(Integer, default=1)

    examples: Mapped[list["DatasetExample"]] = relationship(back_populates="dataset", cascade="all, delete-orphan")


class DatasetExample(Base, UUIDPKMixin, TimestampMixin):
    __tablename__ = "dataset_examples"

    dataset_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("datasets.id"), nullable=False)
    extracted_text_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True), ForeignKey("extracted_texts.id"), nullable=True
    )
    # Lineage: set when this example was produced by Synthetic Expansion (Module 6) from
    # an existing approved example, or by Continuous Learning (Module 15) from a
    # conversation turn.
    parent_example_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True), ForeignKey("dataset_examples.id"), nullable=True
    )
    source_conversation_id: Mapped[uuid.UUID | None] = mapped_column(
        UUID(as_uuid=True), ForeignKey("conversations.id"), nullable=True
    )

    example_type: Mapped[ExampleType] = mapped_column(Enum(ExampleType, name="example_type"), default=ExampleType.qa)
    difficulty: Mapped[Difficulty] = mapped_column(
        Enum(Difficulty, name="difficulty"), default=Difficulty.intermediate
    )

    instruction: Mapped[str] = mapped_column(Text, nullable=False)
    input: Mapped[str | None] = mapped_column(Text, nullable=True)
    output: Mapped[str] = mapped_column(Text, nullable=False)

    status: Mapped[ExampleStatus] = mapped_column(Enum(ExampleStatus, name="example_status"), default=ExampleStatus.pending)
    reviewer_id: Mapped[uuid.UUID | None] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
    review_notes: Mapped[str | None] = mapped_column(Text, nullable=True)
    generated_by: Mapped[str | None] = mapped_column(String(100), nullable=True)

    dataset: Mapped["Dataset"] = relationship(back_populates="examples")