ai-legal-assistant/backend/app/services/contract_service.py

"""Contract service for contract management."""
from typing import List, Optional
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.contract import Contract, ContractTemplate, ContractApproval, ContractStatus


class ContractService:
    """Service for contract operations."""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def create_contract(
        self,
        title: str,
        party_a: str,
        party_b: str,
        content: str,
        created_by: int,
        **kwargs
    ) -> Contract:
        """Create a new contract."""
        contract = Contract(
            title=title,
            party_a=party_a,
            party_b=party_b,
            content=content,
            created_by=created_by,
            **kwargs
        )
        self.db.add(contract)
        await self.db.flush()
        await self.db.refresh(contract)
        return contract

    async def get_contract_by_id(self, contract_id: int) -> Optional[Contract]:
        """Get a contract by ID."""
        result = await self.db.execute(
            select(Contract).where(Contract.id == contract_id)
        )
        return result.scalar_one_or_none()

    async def get_contracts_list(
        self,
        skip: int = 0,
        limit: int = 20,
        status: Optional[ContractStatus] = None,
        created_by: Optional[int] = None,
    ) -> List[Contract]:
        """Get list of contracts."""
        query = select(Contract)

        if status:
            query = query.where(Contract.status == status)
        if created_by:
            query = query.where(Contract.created_by == created_by)

        query = query.offset(skip).limit(limit).order_by(Contract.created_at.desc())
        result = await self.db.execute(query)
        return list(result.scalars().all())

    async def update_contract(
        self,
        contract_id: int,
        **kwargs
    ) -> Optional[Contract]:
        """Update a contract."""
        contract = await self.get_contract_by_id(contract_id)
        if not contract:
            return None

        for key, value in kwargs.items():
            if hasattr(contract, key) and value is not None:
                setattr(contract, key, value)

        await self.db.flush()
        await self.db.refresh(contract)
        return contract

    async def submit_for_approval(self, contract_id: int) -> Optional[Contract]:
        """Submit a contract for approval."""
        return await self.update_contract(
            contract_id,
            status=ContractStatus.PENDING_APPROVAL
        )

    async def approve_contract(
        self,
        contract_id: int,
        approver_id: int,
        comment: Optional[str] = None,
    ) -> ContractApproval:
        """Approve a contract."""
        approval = ContractApproval(
            contract_id=contract_id,
            approver_id=approver_id,
            status=ContractStatus.APPROVED,
            comment=comment,
        )
        self.db.add(approval)

        # Update contract status
        await self.update_contract(contract_id, status=ContractStatus.APPROVED)

        await self.db.flush()
        await self.db.refresh(approval)
        return approval

    async def reject_contract(
        self,
        contract_id: int,
        approver_id: int,
        comment: Optional[str] = None,
    ) -> ContractApproval:
        """Reject a contract."""
        approval = ContractApproval(
            contract_id=contract_id,
            approver_id=approver_id,
            status=ContractStatus.REJECTED,
            comment=comment,
        )
        self.db.add(approval)

        # Update contract status
        await self.update_contract(contract_id, status=ContractStatus.REJECTED)

        await self.db.flush()
        await self.db.refresh(approval)
        return approval


class ContractTemplateService:
    """Service for contract template operations."""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def create_template(
        self,
        name: str,
        content: str,
        created_by: int,
        contract_type: Optional[str] = None,
    ) -> ContractTemplate:
        """Create a contract template."""
        template = ContractTemplate(
            name=name,
            content=content,
            created_by=created_by,
            contract_type=contract_type,
        )
        self.db.add(template)
        await self.db.flush()
        await self.db.refresh(template)
        return template

    async def get_templates_list(
        self,
        skip: int = 0,
        limit: int = 20,
    ) -> List[ContractTemplate]:
        """Get list of templates."""
        result = await self.db.execute(
            select(ContractTemplate)
            .offset(skip)
            .limit(limit)
            .order_by(ContractTemplate.created_at.desc())
        )
        return list(result.scalars().all())