"""
Complete Example: Extending IdentityPlanKit for a SaaS Application

This example shows how to:
1. Extend the User model with additional profile fields
2. Add organization/team management
3. Create custom permissions
4. Integrate everything with your FastAPI app

Production-ready patterns included:
- Unit of Work pattern for transaction safety
- Repository pattern for data access
- Structured logging with context
- Domain exceptions with error codes
- Input validation with Pydantic
- Pagination for list endpoints

See extending_models.html for migration setup instructions.
"""

from datetime import datetime
from typing import Generic, TypeVar
from uuid import UUID

from fastapi import FastAPI, Query, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel as PydanticBaseModel
from pydantic import Field, field_validator
from sqlalchemy import ForeignKey, String, Text, UniqueConstraint, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Mapped, mapped_column, relationship

from identity_plan_kit import CurrentUser, IdentityPlanKit, IdentityPlanKitConfig
from identity_plan_kit.shared.exceptions import BaseError, ConflictError, NotFoundError
from identity_plan_kit.shared.logging import get_logger
from identity_plan_kit.shared.models import BaseModel
from identity_plan_kit.shared.uow import BaseUnitOfWork

logger = get_logger(__name__)

T = TypeVar("T")


# =============================================================================
# MODELS
# =============================================================================


class UserProfile(BaseModel):
    """Extended user profile - separate table linked to IPK users."""

    __tablename__ = "user_profiles"

    user_id: Mapped[UUID] = mapped_column(
        ForeignKey("users.id", ondelete="CASCADE"),
        unique=True,
        index=True,
    )
    phone: Mapped[str | None] = mapped_column(String(20))
    bio: Mapped[str | None] = mapped_column(Text)
    company: Mapped[str | None] = mapped_column(String(255))
    timezone: Mapped[str] = mapped_column(String(50), default="UTC")


class Organization(BaseModel):
    """Organization/team model."""

    __tablename__ = "organizations"

    name: Mapped[str] = mapped_column(String(255))
    slug: Mapped[str] = mapped_column(String(100), unique=True, index=True)
    owner_id: Mapped[UUID] = mapped_column(ForeignKey("users.id"), index=True)

    members: Mapped[list["OrganizationMember"]] = relationship(
        back_populates="organization",
        cascade="all, delete-orphan",
    )


class OrganizationMember(BaseModel):
    """Organization membership."""

    __tablename__ = "organization_members"
    __table_args__ = (
        UniqueConstraint("organization_id", "user_id", name="uq_org_member"),
    )

    organization_id: Mapped[UUID] = mapped_column(
        ForeignKey("organizations.id", ondelete="CASCADE"),
        index=True,
    )
    user_id: Mapped[UUID] = mapped_column(
        ForeignKey("users.id", ondelete="CASCADE"),
        index=True,
    )
    role: Mapped[str] = mapped_column(String(50), default="member")

    organization: Mapped["Organization"] = relationship(back_populates="members")


# =============================================================================
# SCHEMAS
# =============================================================================


class ProfileCreate(PydanticBaseModel):
    """Schema for creating a user profile."""

    phone: str | None = Field(default=None, max_length=20)
    bio: str | None = Field(default=None, max_length=1000)
    company: str | None = Field(default=None, max_length=255)
    timezone: str = Field(default="UTC", max_length=50)


class ProfileResponse(PydanticBaseModel):
    """Schema for profile response."""

    id: UUID
    user_id: UUID
    phone: str | None
    bio: str | None
    company: str | None
    timezone: str
    created_at: datetime

    model_config = {"from_attributes": True}


class OrgCreate(PydanticBaseModel):
    """Schema for creating an organization."""

    name: str = Field(min_length=1, max_length=255)
    slug: str = Field(
        min_length=1,
        max_length=100,
        pattern=r"^[a-z0-9][a-z0-9-]*[a-z0-9]$|^[a-z0-9]$",
        description="URL-safe slug (lowercase alphanumeric and hyphens)",
    )

    @field_validator("slug")
    @classmethod
    def validate_slug(cls, v: str) -> str:
        """Validate slug format."""
        if "--" in v:
            raise ValueError("Slug cannot contain consecutive hyphens")
        return v.lower()


class OrgResponse(PydanticBaseModel):
    """Schema for organization response."""

    id: UUID
    name: str
    slug: str
    owner_id: UUID
    created_at: datetime

    model_config = {"from_attributes": True}


class PaginatedResponse(PydanticBaseModel, Generic[T]):
    """Generic paginated response wrapper."""

    items: list[T]
    total: int
    page: int
    page_size: int
    total_pages: int


# =============================================================================
# REPOSITORIES
# =============================================================================


class ProfileRepository:
    """Repository for UserProfile data access."""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_user_id(self, user_id: UUID) -> UserProfile | None:
        """Get profile by user ID."""
        result = await self._session.execute(
            select(UserProfile).where(UserProfile.user_id == user_id)
        )
        return result.scalar_one_or_none()

    async def create(self, user_id: UUID, **data: str | None) -> UserProfile:
        """Create a new profile."""
        profile = UserProfile(user_id=user_id, **data)
        self._session.add(profile)
        await self._session.flush()
        return profile


class OrganizationRepository:
    """Repository for Organization data access."""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def exists_by_slug(self, slug: str) -> bool:
        """Check if organization with slug exists."""
        result = await self._session.execute(
            select(func.count())
            .select_from(Organization)
            .where(Organization.slug == slug)
        )
        return (result.scalar() or 0) > 0

    async def create(
        self, name: str, slug: str, owner_id: UUID
    ) -> Organization:
        """Create a new organization."""
        org = Organization(name=name, slug=slug, owner_id=owner_id)
        self._session.add(org)
        await self._session.flush()
        return org

    async def list_for_user(
        self, user_id: UUID, page: int, page_size: int
    ) -> tuple[list[Organization], int]:
        """List organizations for a user with pagination."""
        # Count total
        count_query = (
            select(func.count())
            .select_from(Organization)
            .join(OrganizationMember)
            .where(OrganizationMember.user_id == user_id)
        )
        total = (await self._session.execute(count_query)).scalar() or 0

        # Get paginated data
        offset = (page - 1) * page_size
        data_query = (
            select(Organization)
            .join(OrganizationMember)
            .where(OrganizationMember.user_id == user_id)
            .order_by(Organization.created_at.desc())
            .offset(offset)
            .limit(page_size)
        )
        result = await self._session.execute(data_query)
        return list(result.scalars().all()), total


class MemberRepository:
    """Repository for OrganizationMember data access."""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def create(
        self, organization_id: UUID, user_id: UUID, role: str = "member"
    ) -> OrganizationMember:
        """Create a new membership."""
        member = OrganizationMember(
            organization_id=organization_id, user_id=user_id, role=role
        )
        self._session.add(member)
        await self._session.flush()
        return member


# =============================================================================
# UNIT OF WORK
# =============================================================================


class ExtensionUnitOfWork(BaseUnitOfWork):
    """Unit of Work for extension models."""

    profiles: ProfileRepository
    organizations: OrganizationRepository
    members: MemberRepository

    def _init_repositories(self) -> None:
        """Initialize repositories with current session."""
        self.profiles = ProfileRepository(self.session)
        self.organizations = OrganizationRepository(self.session)
        self.members = MemberRepository(self.session)


# =============================================================================
# APPLICATION
# =============================================================================


config = IdentityPlanKitConfig()
kit = IdentityPlanKit(config)

app = FastAPI(
    title="Extended SaaS App",
    description="Example of extending IdentityPlanKit with production patterns",
    lifespan=kit.lifespan,
)
kit.setup(app)


# =============================================================================
# EXCEPTION HANDLERS
# =============================================================================


@app.exception_handler(BaseError)
async def domain_exception_handler(request: Request, exc: BaseError) -> JSONResponse:
    """Handle domain exceptions with structured error responses."""
    logger.warning(
        "domain_error",
        error_code=exc.code,
        message=exc.message,
        path=request.url.path,
        method=request.method,
    )
    return JSONResponse(status_code=exc.status_code, content=exc.to_dict())


# =============================================================================
# PROFILE ENDPOINTS
# =============================================================================


@app.post("/profile", response_model=ProfileResponse)
async def create_profile(data: ProfileCreate, user: CurrentUser) -> UserProfile:
    """
    Create user profile.

    Returns 409 if profile already exists.
    """
    async with ExtensionUnitOfWork(kit.db_manager.session_factory) as uow:
        existing = await uow.profiles.get_by_user_id(user.id)
        if existing:
            raise ConflictError(message="Profile already exists")

        profile = await uow.profiles.create(user_id=user.id, **data.model_dump())

        logger.info(
            "profile_created",
            user_id=str(user.id),
            profile_id=str(profile.id),
        )
        return profile


@app.get("/profile", response_model=ProfileResponse)
async def get_profile(user: CurrentUser) -> UserProfile:
    """
    Get current user's profile.

    Returns 404 if profile not found.
    """
    async with ExtensionUnitOfWork(kit.db_manager.session_factory) as uow:
        profile = await uow.profiles.get_by_user_id(user.id)
        if not profile:
            raise NotFoundError(message="Profile not found")
        return profile


# =============================================================================
# ORGANIZATION ENDPOINTS
# =============================================================================


@app.post("/organizations", response_model=OrgResponse)
async def create_org(data: OrgCreate, user: CurrentUser) -> Organization:
    """
    Create organization. User becomes owner.

    Returns 409 if slug already taken.
    """
    async with ExtensionUnitOfWork(kit.db_manager.session_factory) as uow:
        if await uow.organizations.exists_by_slug(data.slug):
            raise ConflictError(message=f"Slug '{data.slug}' is already taken")

        org = await uow.organizations.create(
            name=data.name,
            slug=data.slug,
            owner_id=user.id,
        )

        # Add owner as member with owner role
        await uow.members.create(
            organization_id=org.id,
            user_id=user.id,
            role="owner",
        )

        logger.info(
            "organization_created",
            org_id=str(org.id),
            slug=data.slug,
            owner_id=str(user.id),
        )
        return org


@app.get("/organizations", response_model=PaginatedResponse[OrgResponse])
async def list_orgs(
    user: CurrentUser,
    page: int = Query(default=1, ge=1, description="Page number"),
    page_size: int = Query(default=20, ge=1, le=100, description="Items per page"),
) -> PaginatedResponse[OrgResponse]:
    """
    List user's organizations with pagination.

    Returns paginated list of organizations the user is a member of.
    """
    async with ExtensionUnitOfWork(kit.db_manager.session_factory) as uow:
        orgs, total = await uow.organizations.list_for_user(user.id, page, page_size)
        total_pages = (total + page_size - 1) // page_size if total > 0 else 0

        logger.debug(
            "organizations_listed",
            user_id=str(user.id),
            page=page,
            total=total,
        )

        return PaginatedResponse(
            items=[OrgResponse.model_validate(org) for org in orgs],
            total=total,
            page=page,
            page_size=page_size,
            total_pages=total_pages,
        )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)
