Template Types

class TemplateType(str, Enum):
    """Template type: RESEARCH (papers), BOOK (books), GENERAL (other)."""
    RESEARCH = "research"
    BOOK = "book"
    GENERAL = "general"

class TemplateParserStep(BasePipelineStep):
    """Analyze template and extract search parameters."""

    async def _execute_step(self, data: PipelineData) -> StepResult:
        agent = create_agent(
            model="anthropic:claude-haiku-4-5",
            output_type=TemplateAnalysis,  # Pydantic model with template_type field
            temperature=0.1,  # Low temperature for consistent classification
            max_tokens=2000
        )

        prompt = f"""
        Analyze this email template and recipient information:

        Template: {data.email_template}
        Recipient: {data.recipient_name}
        Interest: {data.recipient_interest}

        Classify the template type:
        - RESEARCH: Mentions papers, publications, research, studies, or academic work
        - BOOK: Mentions books, textbooks, or published non-academic works
        - GENERAL: Doesn't mention papers or books (networking, opportunities, etc.)

        Extract:
        1. Template type (RESEARCH, BOOK, or GENERAL)
        2. Search terms for finding information about the recipient
        3. Placeholders used ({{name}}, {{research}}, etc.)
        """

        result = await agent.run(prompt)

        # Update PipelineData with classification
        data.template_type = result.data.template_type
        data.search_terms = result.data.search_terms
        data.template_analysis = result.data.model_dump()

        return StepResult(success=True, step_name=self.step_name)

class ArxivEnricherStep(BasePipelineStep):
    """Conditionally fetch ArXiv papers based on template type."""

    async def _execute_step(self, data: PipelineData) -> StepResult:
        # Only execute for RESEARCH templates
        if data.template_type != TemplateType.RESEARCH:
            logfire.info(
                f"Skipping ArXiv enrichment for template type: {data.template_type}",
                task_id=data.task_id
            )
            return StepResult(
                success=True,
                step_name=self.step_name,
                metadata={"papers_fetched": 0, "reason": "non-research template"}
            )

        # Fetch papers for RESEARCH templates
        papers = await self._fetch_arxiv_papers(
            author=data.recipient_name,
            interest=data.recipient_interest
        )

        data.arxiv_papers = papers[:5]  # Top 5 papers
        data.enrichment_metadata = {
            "papers_fetched": len(papers),
            "papers_used": len(data.arxiv_papers)
        }

        return StepResult(success=True, step_name=self.step_name)

class EmailComposerStep(BasePipelineStep):
    """Generate final email with type-aware context."""

    async def _execute_step(self, data: PipelineData) -> StepResult:
        # Build context based on template type
        context = {
            "template": data.email_template,
            "recipient_name": data.recipient_name,
            "recipient_interest": data.recipient_interest,
            "scraped_content": data.scraped_content,
            "template_type": data.template_type,
        }

        # Add papers only for RESEARCH type
        if data.template_type == TemplateType.RESEARCH:
            context["arxiv_papers"] = [
                {
                    "title": p["title"],
                    "abstract": p["abstract"][:500],  # Truncate
                    "year": p["year"],
                }
                for p in data.arxiv_papers[:3]  # Top 3 papers
            ]
            context["instruction"] = "Reference specific paper titles in the email."
        else:
            context["instruction"] = "Use information from web scraping for personalization."

        # Generate email with Sonnet
        email = await self._generate_email(context)

        # Validate with type-aware checks
        validation = self._validate_email(
            email=email,
            template_type=data.template_type,
            has_papers=len(data.arxiv_papers) > 0
        )

        # Write to database
        email_id = await self._write_to_database(data, email)

        return StepResult(success=True, step_name=self.step_name)

class JobStatus(Enum):
    """Pipeline job status for Celery task state management."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"