workflow/financial_workflow.py

"""
Financial Document Analysis Workflow using Agno Workflows
Clean, pure-python implementation with structured outputs to avoid JSON parsing issues
"""

import json
from pathlib import Path
from typing import Dict, List, Optional, Iterator
from pydantic import BaseModel, Field

from agno.agent import Agent, RunResponse
from agno.models.google import Gemini  
from agno.media import File
from agno.tools.file import FileTools
from agno.tools.python import PythonTools
from agno.workflow import Workflow
from agno.utils.log import logger
from agno.tools.shell import ShellTools
from config.settings import settings


# Structured Output Models to avoid JSON parsing issues
class DataPoint(BaseModel):
    """Individual financial data point"""
    field_name: str = Field(..., description="Name of the financial data field")
    value: str = Field(..., description="Value of the field")
    category: str = Field(..., description="Financial category (revenue, expenses, assets, etc.)")
    period: str = Field(default="", description="Time period if applicable")
    unit: str = Field(default="", description="Currency or measurement unit")
    confidence: float = Field(default=0.9, description="Confidence score 0-1")

class ExtractedFinancialData(BaseModel):
    """Structured output for data extraction phase"""
    company_name: str = Field(default="", description="Company name")
    document_type: str = Field(..., description="Type of financial document")
    reporting_period: str = Field(default="", description="Reporting period")
    data_points: List[DataPoint] = Field(..., description="All extracted financial data points")
    summary: str = Field(..., description="Brief summary of extracted data")

class FinancialCategory(BaseModel):
    """A category of organized financial data"""
    category_name: str = Field(..., description="Name of the financial category")
    description: str = Field(..., description="Description of what this category contains")
    data_items: Dict[str, str] = Field(..., description="Key-value pairs of financial data")
    totals: Dict[str, str] = Field(default_factory=dict, description="Any calculated totals")

class ArrangedFinancialData(BaseModel):
    """Structured output for data arrangement phase"""
    categories: List[FinancialCategory] = Field(..., description="Organized financial categories")
    key_metrics: Dict[str, str] = Field(default_factory=dict, description="Key financial metrics")
    insights: List[str] = Field(default_factory=list, description="Financial insights and analysis")
    summary: str = Field(..., description="Summary of arranged data")

class GeneratedCode(BaseModel):
    """Structured output for code generation phase"""
    code: str = Field(..., description="Generated Python code for Excel creation")
    description: str = Field(..., description="Description of what the code does")
    output_filename: str = Field(..., description="Expected output filename")
    execution_notes: str = Field(default="", description="Notes about code execution")


class FinancialDocumentWorkflow(Workflow):
    """
    Pure Python workflow for financial document analysis
    Uses structured outputs to eliminate JSON parsing issues
    """
    
    description: str = "Financial document analysis workflow with data extraction, organization, and Excel generation"
    
    
    
    # Data Extractor Agent - Structured output eliminates JSON parsing issues
    data_extractor: Agent = Agent(
        model=Gemini(id=settings.DATA_EXTRACTOR_MODEL,thinking_budget=settings.DATA_EXTRACTOR_MODEL_THINKING_BUDGET),
        description="Expert financial data extraction specialist",
        instructions=[
                "Extract comprehensive financial data from documents with these priorities:",
                "Identify and classify the document type: Income Statement, Balance Sheet, Cash Flow Statement, 10-K, 10-Q, Annual Report, Quarterly/Interim Report, Prospectus, Earnings Release, Proxy Statement, Investor Presentation, Press Release, or other",
                "Extract report version: audited, unaudited, restated, pro forma",
                "Capture language, country/jurisdiction, and file format (PDF, XLSX, HTML, etc.)",
                "Extract company name and unique identifiers: LEI, CIK, ISIN, Ticker",
                "Extract reporting entity: consolidated, subsidiary, segment",
                "Extract fiscal year and period covered (start and end dates)",
                "Extract all reporting, publication, and filing dates",
                "Extract currency and any currency translation notes",
                "Extract auditors name, if present",
                "Identify financial statement presentation style: single-step, multi-step, consolidated, segmental",
                "Capture table and note references for each data point",
                "Extract total revenue/net sales (with by-product/service, segment, and geography breakdowns if disclosed)",
                "Extract COGS or cost of sales",
                "Extract gross profit and gross margin",
                "Extract operating expenses: R&D, SG&A, advertising, depreciation, amortization",
                "Extract operating income (EBIT) and EBIT margin",
                "Extract non-operating items: interest income/expense, other income/expenses",
                "Extract pretax income, income tax expense, and net income (with breakdowns: continuing, discontinued ops)",
                "Extract basic and diluted EPS",
                "Extract comprehensive and other comprehensive income items",
                "Extract YoY and sequential income comparisons (if available)",
                "Extract current assets: cash and equivalents, marketable securities, accounts receivable (gross/net), inventory (raw, WIP, finished), prepaid expenses, other",
                "Extract non-current assets: PP&E (gross/net), intangible assets, goodwill, LT investments, deferred tax assets, right-of-use assets, other",
                "Extract current liabilities: accounts payable, accrued expenses, short-term debt, lease liabilities, taxes payable, other",
                "Extract non-current liabilities: long-term debt, deferred tax liabilities, pensions, lease obligations, other",
                "Extract total shareholders equity: common/ordinary stock, retained earnings, additional paid-in capital, treasury stock, accumulated OCI, minority interest",
                "Extract book value per share",
                "Extract cash flows: net cash from operating, investing, and financing activities",
                "Extract key cash flow line items: net cash from ops, capex, acquisitions/disposals, dividends, share buybacks, debt activities",
                "Extract non-cash adjustments: depreciation, amortization, SBC, deferred taxes, impairments, gain/loss on sale",
                "Extract profitability ratios: gross margin, operating margin, net margin, EBITDA margin",
                "Extract return ratios: ROE, ROA, ROIC",
                "Extract liquidity/solvency: current ratio, quick ratio, debt/equity, interest coverage",
                "Extract efficiency: asset turnover, inventory turnover, receivables turnover",
                "Extract per-share metrics: EPS (basic/diluted), BVPS, FCF per share",
                "Extract segmental/geographical/operational ratios and breakdowns",
                "Extract shares outstanding, share class details, voting rights",
                "Extract dividends declared/paid (amount, dates)",
                "Extract buyback authorization/utilization details",
                "Extract employee count (average, period-end)",
                "Extract store/branch/office count",
                "Extract customer/user/subscriber numbers (active/paying, ARPU, churn, MAU/DAU)",
                "Extract units shipped/sold, production volumes, operational stats",
                "Extract key management guidance/forecasts if present",
                "Extract risk factors, uncertainties, and forward-looking statements",
                "Extract ESG/sustainability data where available (emissions, board diversity, etc.)",
                "Flag any restatements, adjustments, or one-off items",
                "Highlight material non-recurring, extraordinary, or unusual items (gains/losses, litigation, impairments, restructuring)",
                "Identify related-party transactions and accounting policy changes",
                "For each data point, provide a confidence score (0–1) based on clarity and documentation",
                "Include table/note reference numbers where possible",
                "Note any ambiguity or extraction limitations for specific data",
                "List all units, scales (millions, thousands), and any conversion performed",
                "Normalize date and currency formats across extracted data",
                "Validate calculations (e.g., assets = liabilities + equity), and flag inconsistencies",
                "Return data in a structured format (JSON/table), with reference and confidence annotation"
            ],
        response_model=ExtractedFinancialData,
        structured_outputs=True,
        debug_mode=True,
    )


    
    # Data Arranger Agent - Organizes data into categories for Excel
    data_arranger: Agent = Agent(
        model=Gemini(id=settings.DATA_ARRANGER_MODEL,thinking_budget=settings.DATA_ARRANGER_MODEL_THINKING_BUDGET),
        description="Financial data organization and analysis expert", 
        instructions=[
            'Organize the extracted financial data into logical categories based on financial statement types (Income Statement, Balance Sheet, Cash Flow Statement, etc.).',
            'Group related financial items together (e.g., all revenue items, all expense items, all asset items).',
            'Ensure each category has a clear, descriptive name that would work as an Excel worksheet tab.',
            'Always add appropriate headers for Excel templates including: Years (e.g., 2021, 2022, 2023, 2024), Company names or entity identifiers, Financial line item names, and Units of measurement (e.g., "in millions", "in thousands").',
            'Create column headers that clearly identify what each data column represents.',
            'Include row headers that clearly identify each financial line item.',
            'Design categories suitable for comprehensive Excel worksheets, such as: Income Statement Data, Balance Sheet Data, Cash Flow Data, Key Metrics, and Company Information.',
            'Maintain data integrity - do not modify, calculate, or analyze the original data values.',
            'Preserve original data formats and units.',
            'Ensure data is organized in a tabular format suitable for Excel import.',
            'Include metadata about data sources and reporting periods where available.',
            'Package everything into a JSON object with the fields: categories (object containing organized data by category), headers (object containing appropriate headers for each category), and metadata (object containing information about data sources, periods, and units).',
            'Never perform any analysis on the data.',
            'Do not calculate ratios, growth rates, or trends.',
            'Do not provide insights or interpretations.',
            'Do not modify the actual data values.',
            'Focus solely on organization and proper formatting.',
            'Save this JSON as \'arranged_financial_data.json\' using the save_file tool.',
            'Run list_files to verify that the file now exists in the working directory.',
            'Use read_file to ensure the JSON content was written correctly.',
            'If the file is missing or the content is incorrect, debug, re-save, and repeat steps',
            'Only report success after the files presence and validity are fully confirmed.'
        ],
        tools=[FileTools()],  # FileTools for saving arranged data
        # NOTE: Cannot use structured_outputs with tools in Gemini - choosing tools over structured outputs
        markdown=True,
        debug_mode=True,
        add_memory_references=True,
        add_session_summary_references=True,
        exponential_backoff=True,
        retries=10,
    )
    
    # Code Generator Agent - Creates Excel generation code
    code_generator = Agent(
        model=Gemini(
            id=settings.CODE_GENERATOR_MODEL,
            thinking_budget=settings.CODE_GENERATOR_MODEL_THINKING_BUDGET
        ),
        description="Excel report generator that analyzes JSON data and creates formatted workbooks using shell execution on any OS",
        goal="Generate a professional Excel report from arranged_financial_data.json with multiple worksheets, formatting, and charts",
        instructions=[
            "EXECUTION RULE: Always use run_shell_command() for Python execution. Never use save_to_file_and_run().",
            "",
            "CRITICAL: Always read the file to understand the struction of the JSON First"
            "FIRST, use read_file tool to load 'arranged_financial_data.json'.",
            "SECOND, analyze its structure deeply. Identify all keys, data types, nested structures, and any inconsistencies.",
            "THIRD, create analysis.py to programmatically examine the JSON. Execute using run_shell_command().",
            "FOURTH, based on  the analysis, design your Excel structure. Plan worksheets, formatting, and charts needed.",
            "FIFTH, implement generate_excel_report.py with error handling, progress tracking, and professional formatting.",
            "",
            "CRITICAL: Always start Python scripts with:",
            "import os",
            "os.chdir(os.path.dirname(os.path.abspath(__file__)) or '.')",
            "This ensures the script runs in the correct directory regardless of OS.",
            "",
            "Available Tools:",
            "- FileTools: read_file, save_file, list_files",
            "- PythonTools: pip_install_package (ONLY for package installation)",
            "- ShellTools: run_shell_command (PRIMARY execution tool)",
            "",
            "Cross-Platform Execution:",
            "- Try: run_shell_command('python script.py 2>&1')",
            "- If fails on Windows: run_shell_command('python.exe script.py 2>&1')",
            "- PowerShell alternative: run_shell_command('powershell -Command \"python script.py\" 2>&1')",
            "",
            "Verification Commands (Linux/Mac):",
            "- run_shell_command('ls -la *.xlsx')",
            "- run_shell_command('file Financial_Report*.xlsx')",
            "- run_shell_command('du -h *.xlsx')",
            "",
            "Verification Commands (Windows/PowerShell):",
            "- run_shell_command('dir *.xlsx')",
            "- run_shell_command('powershell -Command \"Get-ChildItem *.xlsx\"')",
            "- run_shell_command('powershell -Command \"(Get-Item *.xlsx).Length\"')",
            "",
            "Debug Commands (Cross-Platform):",
            "- Current directory: run_shell_command('pwd') or run_shell_command('cd')",
            "- Python location: run_shell_command('where python') or run_shell_command('which python')",
            "- List files: run_shell_command('dir') or run_shell_command('ls')",
            "",
            "Package Installation:",
            "- pip_install_package('openpyxl')",
            "- Or via shell: run_shell_command('pip install openpyxl')",
            "- Windows: run_shell_command('python -m pip install openpyxl')",
            "",
            "Success Criteria: Excel file exists, size >5KB, no errors in output."
        ],
        expected_output="A Financial_Report_YYYYMMDD_HHMMSS.xlsx file containing formatted data from the JSON with multiple worksheets, professional styling, and relevant charts",
        additional_context="This agent must work on Windows, Mac, and Linux. Always use os.path for file operations and handle path separators correctly. Include proper error handling for cross-platform compatibility.",
        tools=[
            ShellTools(),
            FileTools(save_files=True, read_files=True, list_files=True),
            PythonTools(pip_install=True, save_and_run=False, run_code=False)
        ],
        markdown=True,
        show_tool_calls=True,
        debug_mode=True,
        retries=10,
        add_datetime_to_instructions=True,
        delay_between_retries=10
            )

    def __init__(self, session_id: str = None, **kwargs):
        super().__init__(session_id=session_id, **kwargs)
        self.session_id = session_id or f"financial_workflow_{int(__import__('time').time())}"
        self.session_output_dir = Path(settings.TEMP_DIR) / self.session_id / "output"
        self.session_output_dir.mkdir(parents=True, exist_ok=True)
        
        # Configure tools with correct base directories after initialization
        self._configure_agent_tools()
        
        logger.info(f"FinancialDocumentWorkflow initialized with session: {self.session_id}")

    def _configure_agent_tools(self):
        """Configure agent tools with the correct base directories"""
        # Configure data arranger's FileTools with session output directory
        if hasattr(self.data_arranger, 'tools') and self.data_arranger.tools:
            for tool in self.data_arranger.tools:
                if isinstance(tool, FileTools):
                    tool.base_dir = self.session_output_dir
                    
        # Configure code generator's tools with session output directory  
        if hasattr(self.code_generator, 'tools') and self.code_generator.tools:
            for tool in self.code_generator.tools:
                if isinstance(tool, FileTools):
                    tool.base_dir = self.session_output_dir
                elif isinstance(tool, PythonTools):
                    tool.base_dir = self.session_output_dir

    def run(self, file_path: str, use_cache: bool = True) -> RunResponse:
        """
        Pure Python workflow execution - no streaming, no JSON parsing issues
        """
        logger.info(f"Processing financial document: {file_path}")
        
        # Check cache first if enabled
        if use_cache and "final_results" in self.session_state:
            logger.info("Returning cached results")
            return RunResponse(
                run_id=self.run_id,
                content=self.session_state["final_results"]
            )
        
        try:
            # Step 1: Extract Financial Data
            logger.info("Step 1: Extracting financial data...")
            
            # Check for cached extraction
            if use_cache and "extracted_data" in self.session_state:
                extracted_data = ExtractedFinancialData.model_validate(
                    self.session_state["extracted_data"]
                )
                logger.info("Using cached extraction data")
            else:
                document = File(filepath=file_path)
                extraction_prompt = f"""
                Analyze this financial document and extract all relevant financial data points.

                Focus on:
                - Company identification, including company name, entity identifiers (e.g., Ticker, CIK, ISIN, LEI), and reporting entity type (consolidated/subsidiary/segment).
                - All reporting period information: fiscal year, period start and end dates, reporting date, publication date, and currency used.
                - Revenue data: total revenue/net sales, breakdown by product/service, segment, and geography if available, and year-over-year growth rates.
                - Expense data: COGS, operating expenses (R&D, SG&A, advertising, depreciation/amortization), interest expenses, taxes, and any non-operating items.
                - Profit data: gross profit, operating income (EBIT/EBITDA), pretax profit, net income, basic and diluted earnings per share (EPS), comprehensive income.
                - Balance sheet items: current assets (cash, securities, receivables, inventories), non-current assets (PP&E, intangibles, goodwill), current liabilities, non-current liabilities, and all categories of shareholders’ equity.
                - Cash flow details: cash from operations, investing, and financing; capex, dividends, buybacks; non-cash adjustments (depreciation, SBC, etc.).
                - Financial ratios: profitability (gross margin, operating margin, net margin), return (ROE, ROA, ROIC), liquidity (current/quick ratio), leverage (debt/equity, interest coverage), efficiency (asset/inventory/receivables turnover), per-share metrics.
                - Capital and shareholder information: shares outstanding, share class details, dividends, and buyback information.
                - Non-financial and operational metrics: employee, store, customer/user counts, production volumes, and operational breakdowns.
                - Extract any additional material metrics, key management guidance, risks, uncertainties, ESG indicators, or forward-looking statements.
                - Flag/annotate any unusual or non-recurring items, restatements, or related-party transactions.
                - For each data point, provide a confidence score (0–1) and, where possible, include reference identifiers (table/note numbers).
                - If units or currencies differ throughout, normalize and annotate the data accordingly.
                Return your extraction in a structured, machine-readable format with references and confidence levels for each field.
                Document path: {file_path}
                """
                
                extraction_response: RunResponse = self.data_extractor.run(
                    extraction_prompt, 
                    files=[document]
                )
                extracted_data: ExtractedFinancialData = extraction_response.content
                
                # Cache the result
                self.session_state["extracted_data"] = extracted_data.model_dump()
                logger.info(f"Extracted {len(extracted_data.data_points)} data points")
            
            # Step 2: Arrange and Organize Data
            logger.info("Step 2: Organizing financial data...")
            
            if use_cache and "arrangement_response" in self.session_state:
                arrangement_content = self.session_state["arrangement_response"]
                logger.info("Using cached arrangement data")
            else:
                arrangement_prompt = f"""
                You are given raw, extracted financial data. Your task is to reorganize it and prepare it for Excel-based reporting.

                ========== WHAT TO DELIVER ==========
                • A single JSON object saved as arranged_financial_data.json
                • Fields required: categories, headers, metadata

                ========== HOW TO ORGANIZE ==========
                Create distinct, Excel-ready categories (one worksheet each) for logical grouping of financial data. Examples include:
                1. Income Statement Data
                2. Balance Sheet Data
                3. Cash Flow Data
                4. Company Information / General Data

                ========== STEP-BY-STEP ==========
                1. Map every data point into the most appropriate category above.
                2. For each category, identify and include all necessary headers for an Excel template, such as years, company names, financial line item names, and units of measurement (e.g., "in millions").
                3. Ensure data integrity by not modifying, calculating, or analyzing the original data values.
                4. Preserve original data formats and units.
                5. Organize data in a tabular format suitable for direct Excel import.
                6. Include metadata about data sources and reporting periods where available.
                7. Assemble everything into the JSON schema described under “WHAT TO DELIVER.”
                8. Save the JSON as arranged_financial_data.json via save_file.
                9. Use list_files to confirm the file exists, then read_file to validate its content.
                10. If the file is missing or malformed, fix the issue and repeat steps 8 – 9.
                11. Only report success after the file passes both existence and content checks.

                ========== IMPORTANT RESTRICTIONS ==========
                - Never perform any analysis on the data.
                - Do not calculate ratios, growth rates, or trends.
                - Do not provide insights or interpretations.
                - Do not modify the actual data values.
                - Focus solely on organization and proper formatting.
                
                Extracted Data: {extracted_data.model_dump_json(indent=2)}
                """
                
                arrangement_response: RunResponse = self.data_arranger.run(arrangement_prompt)
                arrangement_content = arrangement_response.content
                
                # Cache the result
                self.session_state["arrangement_response"] = arrangement_content
                logger.info("Data organization completed - check output directory for arranged_financial_data.json")
            
            # Step 3: Generate and Execute Excel Code
            logger.info("Step 3: Generating and executing Excel code...")
            
            if use_cache and "code_generation_response" in self.session_state:
                code_generation_content = self.session_state["code_generation_response"]
                execution_success = self.session_state.get("execution_success", False)
                logger.info("Using cached code generation results")
            else:
                code_prompt = f"""
                Your objective: Turn the organized JSON data into a polished, multi-sheet Excel report—and prove that it works.

                ========== INPUT ==========
                File: arranged_financial_data.json  
                Tool to read it: read_file

                ========== WHAT THE PYTHON SCRIPT MUST DO ==========
                1. Load arranged_financial_data.json and parse its contents.
                2. For each category in the JSON, create a dedicated worksheet using openpyxl.
                3. Apply professional touches:
                • Bold, centered headers  
                • Appropriate number formats  
                • Column-width auto-sizing  
                • Borders, cell styles, and freeze panes
                4. Insert charts (bar, line, or pie) wherever the data lends itself to visualisation.
                5. Embed key metrics and summary notes prominently in the Executive Summary sheet.
                6. Name the workbook: Financial_Report_<YYYYMMDD_HHMMSS>.xlsx.
                7. Wrap every file and workbook operation in robust try/except blocks.
                8. Log all major steps and any exceptions for easy debugging.
                9. Save the script via save_to_file_and_run and execute it immediately.
                10. After execution, use list_files to ensure the Excel file was created.
                11. Optionally inspect the file (e.g., size or first bytes via read_file) to confirm it is not empty.
                12. If the workbook is missing or corrupted, refine the code, re-save, and re-run until success.

                ========== OUTPUT ==========
                • A fully formatted Excel workbook in the working directory.
                • A concise summary of what ran, any issues encountered, and confirmation that the file exists and opens without error.
                """
                
                code_response: RunResponse = self.code_generator.run(code_prompt)
                code_generation_content = code_response.content
                
                # Simple check for execution success based on response content
                execution_success = (
                    "error" not in code_generation_content.lower() or
                    "success" in code_generation_content.lower() or
                    "completed" in code_generation_content.lower()
                )
                
                # Cache the results
                self.session_state["code_generation_response"] = code_generation_content
                self.session_state["execution_success"] = execution_success
                
                logger.info(f"Code generation and execution completed: {'✅ Success' if execution_success else '❌ Failed'}")
            
            # Prepare final results
            # List actual output files
            output_files = []
            if self.session_output_dir.exists():
                output_files = [f.name for f in self.session_output_dir.iterdir() if f.is_file()]
            
            results_summary = f"""
# Financial Document Analysis Complete

## Document Information
- **Company**: {extracted_data.company_name or 'Not specified'}
- **Document Type**: {extracted_data.document_type}
- **Reporting Period**: {extracted_data.reporting_period or 'Not specified'}

## Processing Summary
- **Data Points Extracted**: {len(extracted_data.data_points)}
- **Data Organization**: {'✅ Completed' if arrangement_content else '❌ Failed'}
- **Excel Creation**: {'✅ Success' if execution_success else '❌ Failed'}

## Data Organization Results
{arrangement_content[:500] + '...' if arrangement_content and len(arrangement_content) > 500 else arrangement_content or 'No arrangement data available'}

## Tool Execution Summary
**Data Arranger**: Used FileTools to save organized data to JSON
**Code Generator**: Used PythonTools and FileTools for Excel generation

## Code Generation Results  
{code_generation_content[:500] + '...' if code_generation_content and len(code_generation_content) > 500 else code_generation_content or 'No code generation results available'}

## Generated Files ({len(output_files)} files)
{chr(10).join(f"- **{file}**" for file in output_files) if output_files else "- No files generated"}

## Output Directory
📁 `{self.session_output_dir}`

---
*Generated using Agno Workflows with FileTools and PythonTools integration*
*Note: Due to Gemini limitations, structured outputs were used for data extraction only*
            """
            
            # Cache final results
            self.session_state["final_results"] = results_summary
            
            return RunResponse(
                run_id=self.run_id,
                content=results_summary
            )
            
        except Exception as e:
            error_message = f"❌ Workflow failed: {str(e)}"
            logger.error(f"Financial workflow error: {e}", exc_info=True)
            return RunResponse(
                run_id=self.run_id,
                content=error_message
            )

    def get_processing_status(self) -> Dict[str, str]:
        """Get the current processing status"""
        status = {
            "extraction": "completed" if "extracted_data" in self.session_state else "pending",
            "arrangement": "completed" if "arranged_data" in self.session_state else "pending", 
            "code_generation": "completed" if "generated_code" in self.session_state else "pending",
            "final_results": "completed" if "final_results" in self.session_state else "pending"
        }
        return status