ai-legal-assistant/backend/app/api/v1/phase2.py

"""Phase 2 API endpoints."""
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import get_db
from app.models.phase2 import EntityType, RelationType, RiskLevel, RiskType
from app.schemas.phase2 import (
    RiskAnalysisRequest,
    RiskResponse,
    PredictionRequest,
    PredictionResponse,
    EntityCreate,
    EntityResponse,
    RelationCreate,
    RelationResponse,
    TranslationRequest,
    TranslationResponse,
    LawyerCreate,
    LawyerResponse,
    RecommendationResponse,
)
from app.services.phase2_services import (
    RiskAnalysisService,
    CasePredictionService,
    KnowledgeGraphService,
    TranslationService,
    LawyerMatchingService,
)

router = APIRouter(tags=["phase2"])


# ============ Risk Analysis ============
@router.post("/risks/analyze", response_model=list[RiskResponse])
async def analyze_risks(
    request: RiskAnalysisRequest,
    db: AsyncSession = Depends(get_db),
):
    """Analyze contract for risks."""
    service = RiskAnalysisService(db)
    risks = await service.analyze_contract_risks(
        contract_id=request.contract_id,
        contract_content=request.contract_content,
    )
    return risks


@router.get("/risks/{contract_id}", response_model=list[RiskResponse])
async def get_risks(
    contract_id: int,
    db: AsyncSession = Depends(get_db),
):
    """Get risks for a contract."""
    service = RiskAnalysisService(db)
    risks = await service.get_risks_by_contract(contract_id)
    return risks


# ============ Case Prediction ============
@router.post("/predictions", response_model=PredictionResponse)
async def predict_case(
    request: PredictionRequest,
    user_id: int = 1,  # TODO: Get from auth
    db: AsyncSession = Depends(get_db),
):
    """Predict case outcome."""
    service = CasePredictionService(db)
    prediction = await service.predict_case(
        user_id=user_id,
        case_description=request.case_description,
    )
    return prediction


# ============ Knowledge Graph ============
@router.post("/knowledge-graph/entities", response_model=EntityResponse)
async def create_entity(
    data: EntityCreate,
    db: AsyncSession = Depends(get_db),
):
    """Create a knowledge graph entity."""
    service = KnowledgeGraphService(db)
    entity = await service.create_entity(
        name=data.name,
        entity_type=EntityType(data.entity_type),
        properties=data.properties,
    )
    return entity


@router.get("/knowledge-graph/entities", response_model=list[EntityResponse])
async def search_entities(
    keyword: str = Query(..., min_length=1),
    limit: int = Query(20, ge=1, le=100),
    db: AsyncSession = Depends(get_db),
):
    """Search knowledge graph entities."""
    service = KnowledgeGraphService(db)
    entities = await service.search_entities(keyword, limit)
    return entities


@router.post("/knowledge-graph/relations", response_model=RelationResponse)
async def create_relation(
    data: RelationCreate,
    db: AsyncSession = Depends(get_db),
):
    """Create a relation between entities."""
    service = KnowledgeGraphService(db)
    relation = await service.create_relation(
        source_id=data.source_id,
        target_id=data.target_id,
        relation_type=RelationType(data.relation_type),
        weight=data.weight,
    )
    return relation


# ============ Translation ============
@router.post("/translations", response_model=TranslationResponse)
async def create_translation(
    data: TranslationRequest,
    db: AsyncSession = Depends(get_db),
):
    """Create and perform translation."""
    service = TranslationService(db)
    record = await service.create_translation_request(
        source_text=data.source_text,
        source_lang=data.source_lang,
        target_lang=data.target_lang,
    )
    # Perform translation
    result = await service.translate(record.id)
    return result


# ============ Lawyer Matching ============
@router.post("/lawyers", response_model=LawyerResponse)
async def create_lawyer(
    data: LawyerCreate,
    db: AsyncSession = Depends(get_db),
):
    """Create a lawyer profile."""
    service = LawyerMatchingService(db)
    lawyer = await service.create_lawyer(
        name=data.name,
        specialties=data.specialties,
        experience_years=data.experience_years,
    )
    return lawyer


@router.get("/lawyers", response_model=list[LawyerResponse])
async def list_lawyers(
    specialty: Optional[str] = Query(None),
    limit: int = Query(20, ge=1, le=100),
    db: AsyncSession = Depends(get_db),
):
    """List lawyers."""
    service = LawyerMatchingService(db)
    lawyers = await service.get_lawyers_list(specialty, limit)
    return lawyers


@router.post("/lawyers/recommend", response_model=list[RecommendationResponse])
async def recommend_lawyers(
    case_description: str = Query(..., min_length=10),
    limit: int = Query(5, ge=1, le=20),
    db: AsyncSession = Depends(get_db),
):
    """Recommend lawyers for a case."""
    service = LawyerMatchingService(db)
    recommendations = await service.recommend_lawyers(case_description, limit)
    return recommendations