app/routers/chat.py

from datetime import datetime
from typing import Annotated, Literal

from fastapi import APIRouter, HTTPException
from pydantic import BaseModel, ConfigDict, Field
from app.websockets.handlers import get_graph_manager
from app.services.citation_graph import CitationGraphService
from app.services.paper_acquisition import PaperAcquisitionResult, PaperAcquisitionService
from app.services.project_activity import ProjectActivityService
from app.services.recommendation_clarification import (
    ClarificationExpiredError,
    ClarificationProjectError,
    ClarificationStateError,
    RecommendationClarification,
    RecommendationOutcome,
    RecommendationResult,
)
from app.services.recommendation_service import RecommendationService
from app.services.scholar_service import fetch_top_papers
from app.services.paper_metadata import PaperAuthor, PaperDestination

router = APIRouter()
_activity = ProjectActivityService()

class RecommendRequest(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    project_id: str = Field(min_length=1, max_length=128)
    query: str | None = Field(default=None, max_length=4000)
    project_context_enabled: bool = True
    selection_text: str | None = Field(default=None, max_length=4000)

class PaperResponse(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)

    candidate_id: str = ""
    recommendation_id: str = ""
    title: str
    authors: str
    year: int | None
    cited_by: int
    url: str
    abstract: str = ""
    why_recommended: str = ""
    based_on: str = ""
    source: str = ""
    venue: str = ""
    doi: str = ""
    arxiv_id: str = ""
    is_open_access: bool = False
    pdf_url: str = ""
    provider: str = ""
    pdf_provider: str = ""
    field_providers: dict[str, str] = Field(default_factory=dict)
    resolution_confidence: float = 0.0
    semantic_scholar_id: str = ""
    openalex_id: str = ""
    pmid: str = ""
    openreview_id: str = ""
    author_details: list[PaperAuthor] = Field(default_factory=list)
    reference_count: int = 0
    destinations: list[PaperDestination] = Field(default_factory=list)
    metadata_quality: str = "sparse"
    abstract_status: str = ""
    looked_up_at: float = 0.0
    error: str = ""
    created_at: float = 0.0
    can_acquire: bool = False
    discovery_intent: str = "general"
    discovery_modes: list[str] = Field(default_factory=list)
    discovery_rank: int = 0
    year_bucket: int | None = None
    year_bucket_rank: float | None = None
    rank_factors: dict[str, float] = Field(default_factory=dict)
    matched_queries: list[str] = Field(default_factory=list)
    seed_titles: list[str] = Field(default_factory=list)
    domain_anchors: list[str] = Field(default_factory=list)
    search_queries: list[str] = Field(default_factory=list)
    seed_relation: str = ""
    seed_relative_score: float = 0.0
    anchor_hits: list[str] = Field(default_factory=list)
    citation_verified: bool = False
    verified_citation_direction: str = ""
    verified_seed_openalex_id: str = ""


class RecommendationCoverageResponse(BaseModel):
    requested: int
    returned: int
    complete: bool
    message: str | None = None


class ClarificationChoiceResponse(BaseModel):
    id: str
    label: str


class RecommendationResponse(BaseModel):
    kind: Literal["recommendations"] = "recommendations"
    papers: list[PaperResponse]
    assumptions: list[str] = Field(default_factory=list)
    coverage: RecommendationCoverageResponse


class ClarificationRequiredResponse(BaseModel):
    kind: Literal["clarification_required"] = "clarification_required"
    clarification_token: str
    question: str
    choices: list[ClarificationChoiceResponse]
    expires_at: datetime


RecommendationHttpOutcome = Annotated[
    RecommendationResponse | ClarificationRequiredResponse,
    Field(discriminator="kind"),
]


class RecommendationContinuationRequest(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    project_id: str = Field(min_length=1, max_length=128)
    clarification_token: str = Field(min_length=16, max_length=256)
    answer: str = Field(min_length=1, max_length=1000)
    selection_text: str | None = Field(default=None, max_length=4000)


class AddRecommendationRequest(BaseModel):
    model_config = ConfigDict(extra="forbid", strict=True)
    project_id: str = Field(min_length=1, max_length=128)

def _recommendation_context(project_id: str, enabled: bool):
    if enabled:
        graph = get_graph_manager()
        nodes = graph.list_nodes(project_id)
        labels = [node.label for node in sorted(nodes, key=lambda item: item.depth)[:4]]
        return CitationGraphService().load(project_id), labels
    from app.services.citation_graph import CitationGraph
    return CitationGraph(project_id=project_id), []


def _http_outcome(outcome: RecommendationOutcome) -> RecommendationHttpOutcome:
    if isinstance(outcome, RecommendationClarification):
        return ClarificationRequiredResponse(
            clarification_token=outcome.clarification_token,
            question=outcome.question,
            choices=[ClarificationChoiceResponse(id=choice.id, label=choice.label) for choice in outcome.choices],
            expires_at=outcome.expires_at,
        )
    return RecommendationResponse(
        papers=[PaperResponse.model_validate(paper) for paper in outcome.papers],
        assumptions=outcome.assumptions,
        coverage=RecommendationCoverageResponse.model_validate(outcome.coverage.model_dump()),
    )


@router.post("/recommend", response_model=RecommendationHttpOutcome)
async def recommend_papers(req: RecommendRequest):
    if not CitationGraphService.is_valid_project_id(req.project_id):
        raise HTTPException(400, "Invalid project id")
    citation_graph, graph_labels = _recommendation_context(req.project_id, req.project_context_enabled)
    service = RecommendationService(fetcher=fetch_top_papers)
    _activity.record(req.project_id, "recommendation", "Paper recommendation planning", status="running", metadata={"query": req.query or ""})
    outcome = await service.recommend_or_clarify(
        project_id=req.project_id,
        query=req.query or "",
        citation_graph=citation_graph,
        graph_labels=graph_labels,
        use_project_context=req.project_context_enabled,
        selection_text=req.selection_text,
    )
    if isinstance(outcome, RecommendationClarification):
        _activity.record(req.project_id, "recommendation", "Paper discovery is awaiting one clarification", metadata={"question": outcome.question})
    else:
        _activity.record(req.project_id, "recommendation", "Paper recommendations ready", metadata={"count": len(outcome.papers), "complete": outcome.coverage.complete})
    return _http_outcome(outcome)


@router.post("/recommend/clarify", response_model=RecommendationHttpOutcome)
async def continue_recommendation(req: RecommendationContinuationRequest):
    if not CitationGraphService.is_valid_project_id(req.project_id):
        raise HTTPException(400, "Invalid project id")
    service = RecommendationService(fetcher=fetch_top_papers)
    try:
        # Project-context preference is persisted in the pending token and applied by the service.
        citation_graph, graph_labels = _recommendation_context(req.project_id, True)
        outcome: RecommendationResult = await service.continue_after_clarification(
            project_id=req.project_id,
            token=req.clarification_token,
            answer=req.answer,
            citation_graph=citation_graph,
            graph_labels=graph_labels,
            selection_text=req.selection_text,
        )
    except ClarificationExpiredError as exc:
        raise HTTPException(410, str(exc)) from exc
    except ClarificationProjectError as exc:
        raise HTTPException(403, str(exc)) from exc
    except ClarificationStateError as exc:
        raise HTTPException(404, str(exc)) from exc
    _activity.record(req.project_id, "recommendation", "Paper discovery resumed after clarification", metadata={"count": len(outcome.papers), "complete": outcome.coverage.complete})
    return _http_outcome(outcome)


@router.post("/recommendations/{recommendation_id}/add-to-project", response_model=PaperAcquisitionResult)
async def add_recommendation_to_project(recommendation_id: str, req: AddRecommendationRequest):
    if not CitationGraphService.is_valid_project_id(req.project_id):
        raise HTTPException(400, "Invalid project id")
    paper = RecommendationService(fetcher=fetch_top_papers).get_cached(req.project_id, recommendation_id)
    if paper is None:
        raise HTTPException(404, "Recommendation not found")
    return await PaperAcquisitionService().acquire(req.project_id, paper)