"""EGRUL Workflow - Get company information from Russian EGRUL registry."""

from pydantic import BaseModel, Field

from agno.workflow import Workflow, RunResponse


class EgrulInput(BaseModel):
    """Input for EGRUL workflow."""

    inn: str = Field(..., description="Company INN (10 or 12 digits)")


class EgrulOutput(BaseModel):
    """Output of EGRUL workflow."""

    inn: str
    success: bool
    company_name: str | None = None
    ogrn: str | None = None
    pdf_path: str | None = None
    error: str | None = None


class EgrulWorkflow(Workflow):
    """Workflow for retrieving company data from Russian EGRUL registry.

    This workflow:
    1. Takes an INN (company tax ID) as input
    2. Scrapes egrul.nalog.ru to find the company
    3. Downloads the official PDF extract
    4. Returns company name, OGRN, and PDF path
    """

    name: str = "EGRUL Workflow"
    description: str = "Get company information from Russian EGRUL registry by INN"

    async def arun(self, inn: str = None, **kwargs) -> RunResponse:
        """Execute the EGRUL scraping workflow."""
        from tools.egrul import scrape_egrul

        # Support both direct call and Playground API call
        if inn is None:
            inn = kwargs.get("inn", "")

        # Run the scraper
        result = await scrape_egrul(inn)

        # Build structured output
        output = EgrulOutput(
            inn=result.inn,
            success=result.success,
            company_name=result.company_name,
            ogrn=result.ogrn,
            pdf_path=result.pdf_path,
            error=result.error,
        )

        # Format response content
        if result.success:
            content = f"""Company found!

**INN:** {result.inn}
**Company Name:** {result.company_name or 'N/A'}
**OGRN:** {result.ogrn or 'N/A'}
**PDF saved to:** {result.pdf_path}"""
        else:
            content = f"""Failed to retrieve company data.

**INN:** {result.inn}
**Error:** {result.error}"""

        return RunResponse(
            content=content,
        )

    def run(self, inn: str = None, **kwargs) -> RunResponse:
        """Synchronous wrapper for the workflow."""
        import asyncio
        import nest_asyncio

        # Allow nested event loops (needed when called from FastAPI)
        nest_asyncio.apply()

        # Support both direct call and Playground API call
        if inn is None:
            inn = kwargs.get("inn", "")

        return asyncio.get_event_loop().run_until_complete(self.arun(inn=inn, **kwargs))

