# pycorpfin/tools.py
"""AI Agent tool wrappers for pyCorpFin"""

from typing import List, Dict, Union
from datetime import datetime
from .core import (
    IRR,
    calculate_mirr,
    calculate_xirr,
    presentValue,
    futureValue,
    emi,
    annuity
)

try:
    from langchain_core.tools import StructuredTool
    from pydantic import BaseModel, Field

    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False

# Pydantic models for structured inputs
if LANGCHAIN_AVAILABLE:
    class IRRInput(BaseModel):
        cashflows: List[float] = Field(
            description="List of cash flows (must contain at least one positive and one negative value). Example: [-1000, 300, 400, 500]"
        )
        guess: float = Field(
            default=0.1,
            description="Initial guess for IRR (default 0.1 for 10%)"
        )


    class MIRRInput(BaseModel):
        cashflows: List[float] = Field(
            description="List of cash flows. Example: [-1000, 300, 400, 500]"
        )
        finance_rate: float = Field(
            description="Financing rate as decimal (e.g., 0.10 for 10%)"
        )
        reinvest_rate: float = Field(
            description="Reinvestment rate as decimal (e.g., 0.08 for 8%)"
        )


    class XIRRInput(BaseModel):
        cashflows: List[float] = Field(
            description="List of cash flows. Example: [-1000, 300, 400, 500]"
        )
        dates: List[str] = Field(
            description="List of dates in YYYY-MM-DD format. Example: ['2023-01-01', '2023-06-01', '2023-12-01', '2024-06-01']"
        )
        guess: float = Field(
            default=0.1,
            description="Initial guess for XIRR (default 0.1 for 10%)"
        )


    class PresentValueInput(BaseModel):
        cashflows: List[float] = Field(
            description="List of cash flows for each period. Example: [1000, 500, 300, 200]"
        )
        discount_rate: float = Field(
            description="Discount rate as decimal (e.g., 0.10 for 10%)"
        )


    class FutureValueInput(BaseModel):
        cashflows: List[float] = Field(
            description="List of cash flows for each period. Example: [1000, 500, 300, 200]"
        )
        interest_rate: float = Field(
            description="Interest rate as decimal (e.g., 0.10 for 10%)"
        )


    class EMIInput(BaseModel):
        principal: float = Field(
            description="Principal loan amount in currency units. Example: 100000"
        )
        annual_rate: float = Field(
            description="Annual interest rate as decimal (e.g., 0.08 for 8%)"
        )
        months: int = Field(
            description="Loan tenure in months. Example: 240 for 20 years"
        )


    class AnnuityInput(BaseModel):
        present_value: float = Field(
            description="Present value or principal amount"
        )
        rate: float = Field(
            description="Interest rate per period as decimal (e.g., 0.01 for 1%)"
        )
        periods: int = Field(
            description="Number of periods"
        )


def irr_tool_func(cashflows: List[float], guess: float = 0.1) -> str:
    """Calculate Internal Rate of Return for a series of cash flows."""
    result = IRR(cashflows, guess)
    if isinstance(result, str):
        return f"Error: {result}"
    return f"IRR: {result:.4%} (or {result:.6f} as decimal)"


def mirr_tool_func(cashflows: List[float], finance_rate: float, reinvest_rate: float) -> str:
    """Calculate Modified Internal Rate of Return."""
    result = calculate_mirr(cashflows, finance_rate, reinvest_rate)
    if isinstance(result, str):
        return f"Error: {result}"
    return f"MIRR: {result:.4%} (or {result:.6f} as decimal)"


def xirr_tool_func(cashflows: List[float], dates: List[str], guess: float = 0.1) -> str:
    """Calculate Extended Internal Rate of Return for irregular cash flows with specific dates."""
    try:
        date_objects = [datetime.strptime(d, "%Y-%m-%d") for d in dates]
    except ValueError as e:
        return f"Error: Invalid date format. Use YYYY-MM-DD. Details: {str(e)}"

    result = calculate_xirr(cashflows, date_objects, guess)
    if isinstance(result, str):
        return f"Error: {result}"
    return f"XIRR: {result:.4%} (or {result:.6f} as decimal)"


def present_value_tool_func(cashflows: List[float], discount_rate: float) -> str:
    """Calculate Net Present Value of a series of cash flows."""
    n = len(cashflows) - 1
    cf_dict = {i: cashflows[i] for i in range(len(cashflows))}
    result = presentValue(cf_dict, discount_rate, n)
    return f"Present Value (NPV): ${result:,.2f}"


def future_value_tool_func(cashflows: List[float], interest_rate: float) -> str:
    """Calculate Future Value of a series of cash flows."""
    n = len(cashflows) - 1
    cf_dict = {i: cashflows[i] for i in range(len(cashflows))}
    result = futureValue(cf_dict, interest_rate, n)
    return f"Future Value: ${result:,.2f}"


def emi_tool_func(principal: float, annual_rate: float, months: int) -> str:
    """Calculate Equated Monthly Installment for a loan."""
    monthly_rate = annual_rate / 12
    result = emi(principal, monthly_rate, months)
    total_payment = result * months
    total_interest = total_payment - principal

    return (f"EMI: ${result:,.2f} per month\n"
            f"Total Payment: ${total_payment:,.2f}\n"
            f"Total Interest: ${total_interest:,.2f}\n"
            f"Principal: ${principal:,.2f}")


def annuity_tool_func(present_value: float, rate: float, periods: int) -> str:
    """Calculate annuity payment given present value, rate, and number of periods."""
    result = annuity(present_value, rate, periods)
    total_payment = result * periods
    total_interest = total_payment - present_value

    return (f"Annuity Payment: ${result:,.2f} per period\n"
            f"Total Payment: ${total_payment:,.2f}\n"
            f"Total Interest: ${total_interest:,.2f}")


# LangChain Tools
def get_langchain_tools():
    """Get all pyCorpFin tools as LangChain StructuredTools."""
    if not LANGCHAIN_AVAILABLE:
        raise ImportError(
            "LangChain is not installed. Install with: pip install pycorpfin[langchain]"
        )

    return [
        StructuredTool.from_function(
            func=irr_tool_func,
            name="calculate_irr",
            description=(
                "Calculate Internal Rate of Return (IRR) for a series of cash flows. "
                "IRR is the discount rate that makes the net present value (NPV) of all cash flows equal to zero. "
                "Use this when you need to find the rate of return for an investment with periodic cash flows."
            ),
            args_schema=IRRInput
        ),
        StructuredTool.from_function(
            func=mirr_tool_func,
            name="calculate_mirr",
            description=(
                "Calculate Modified Internal Rate of Return (MIRR) with separate financing and reinvestment rates. "
                "MIRR addresses some limitations of IRR by assuming positive cash flows are reinvested at the "
                "reinvestment rate and negative cash flows are financed at the financing rate."
            ),
            args_schema=MIRRInput
        ),
        StructuredTool.from_function(
            func=xirr_tool_func,
            name="calculate_xirr",
            description=(
                "Calculate Extended Internal Rate of Return (XIRR) for irregular cash flows occurring on specific dates. "
                "Use this instead of IRR when cash flows don't occur at regular intervals. "
                "Dates must be provided in YYYY-MM-DD format."
            ),
            args_schema=XIRRInput
        ),
        StructuredTool.from_function(
            func=present_value_tool_func,
            name="calculate_present_value",
            description=(
                "Calculate Present Value (NPV) of a series of future cash flows discounted at a given rate. "
                "This tells you what future cash flows are worth in today's dollars."
            ),
            args_schema=PresentValueInput
        ),
        StructuredTool.from_function(
            func=future_value_tool_func,
            name="calculate_future_value",
            description=(
                "Calculate Future Value of a series of cash flows compounded at a given interest rate. "
                "This tells you what current cash flows will be worth at a future date."
            ),
            args_schema=FutureValueInput
        ),
        StructuredTool.from_function(
            func=emi_tool_func,
            name="calculate_emi",
            description=(
                "Calculate Equated Monthly Installment (EMI) for a loan given principal, annual interest rate, and tenure. "
                "Returns monthly payment amount, total payment, and total interest. Common for mortgages and personal loans."
            ),
            args_schema=EMIInput
        ),
        StructuredTool.from_function(
            func=annuity_tool_func,
            name="calculate_annuity",
            description=(
                "Calculate periodic annuity payment given present value, interest rate per period, and number of periods. "
                "Similar to EMI but more general - can be used for any periodic payment calculation."
            ),
            args_schema=AnnuityInput
        )
    ]


# Simple dictionary-based tools for frameworks without LangChain
def get_generic_tools():
    """Get tools as simple dictionaries for any AI framework."""
    return [
        {
            "name": "calculate_irr",
            "description": "Calculate Internal Rate of Return (IRR) for cash flows",
            "parameters": {
                "type": "object",
                "properties": {
                    "cashflows": {
                        "type": "array",
                        "items": {"type": "number"},
                        "description": "List of cash flows"
                    },
                    "guess": {
                        "type": "number",
                        "default": 0.1,
                        "description": "Initial guess for IRR"
                    }
                },
                "required": ["cashflows"]
            },
            "function": irr_tool_func
        },
        {
            "name": "calculate_mirr",
            "description": "Calculate Modified Internal Rate of Return (MIRR)",
            "parameters": {
                "type": "object",
                "properties": {
                    "cashflows": {
                        "type": "array",
                        "items": {"type": "number"}
                    },
                    "finance_rate": {"type": "number"},
                    "reinvest_rate": {"type": "number"}
                },
                "required": ["cashflows", "finance_rate", "reinvest_rate"]
            },
            "function": mirr_tool_func
        },
        {
            "name": "calculate_xirr",
            "description": "Calculate Extended Internal Rate of Return (XIRR)",
            "parameters": {
                "type": "object",
                "properties": {
                    "cashflows": {
                        "type": "array",
                        "items": {"type": "number"}
                    },
                    "dates": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Dates in YYYY-MM-DD format"
                    },
                    "guess": {
                        "type": "number",
                        "default": 0.1
                    }
                },
                "required": ["cashflows", "dates"]
            },
            "function": xirr_tool_func
        },
        {
            "name": "calculate_present_value",
            "description": "Calculate Present Value (NPV) of cash flows",
            "parameters": {
                "type": "object",
                "properties": {
                    "cashflows": {
                        "type": "array",
                        "items": {"type": "number"}
                    },
                    "discount_rate": {"type": "number"}
                },
                "required": ["cashflows", "discount_rate"]
            },
            "function": present_value_tool_func
        },
        {
            "name": "calculate_future_value",
            "description": "Calculate Future Value of cash flows",
            "parameters": {
                "type": "object",
                "properties": {
                    "cashflows": {
                        "type": "array",
                        "items": {"type": "number"}
                    },
                    "interest_rate": {"type": "number"}
                },
                "required": ["cashflows", "interest_rate"]
            },
            "function": future_value_tool_func
        },
        {
            "name": "calculate_emi",
            "description": "Calculate Equated Monthly Installment for a loan",
            "parameters": {
                "type": "object",
                "properties": {
                    "principal": {"type": "number"},
                    "annual_rate": {"type": "number"},
                    "months": {"type": "integer"}
                },
                "required": ["principal", "annual_rate", "months"]
            },
            "function": emi_tool_func
        },
        {
            "name": "calculate_annuity",
            "description": "Calculate periodic annuity payment",
            "parameters": {
                "type": "object",
                "properties": {
                    "present_value": {"type": "number"},
                    "rate": {"type": "number"},
                    "periods": {"type": "integer"}
                },
                "required": ["present_value", "rate", "periods"]
            },
            "function": annuity_tool_func
        }
    ]