import requests
import gradio as gr
from langchain_core.tools import tool
from typing import Annotated, Optional

FINANCIAL_DATASETS_API_BASE = "https://api.financialdatasets.ai"

STATEMENT_TYPE_KEY_MAP = {
    "income-statements": "income_statements",
    "balance-sheets": "balance_sheets",
    "cash-flow-statements": "cash_flow_statements"
}

# --- LangChain Tool ---
@tool
def get_financial_statements(
    ticker: Annotated[str, "e.g. AAPL"],
    report_period_lte: Annotated[str, "End date (YYYY-MM-DD)"],
    report_period_gte: Annotated[str, "Start date (YYYY-MM-DD)"],
    statement_type: Annotated[Optional[str], "income-statements, balance-sheets, cash-flow-statements, or 'all'"] = None,
    period: Annotated[str, "annual, quarterly, or ttm"] = "annual",
    limit: Annotated[int, "1-10 records"] = 4,
    api_key: Annotated[str, "Your FinancialDatasets API Key"] = ""
) -> str:
    """Fetch financial statements."""
    if not api_key:
        return "❌ Error: API key is required."

    headers = {"X-API-KEY": api_key}

    if not statement_type or statement_type == "all":
        url = (
            f"{FINANCIAL_DATASETS_API_BASE}/financials"
            f"?ticker={ticker}&period={period}&limit={limit}"
            f"&report_period_lte={report_period_lte}&report_period_gte={report_period_gte}"
        )
        response = requests.get(url, headers=headers)
        if response.status_code != 200:
            return f"❌ API Error: {response.status_code} - {response.text}"
        return response.text

    json_key = STATEMENT_TYPE_KEY_MAP.get(statement_type)
    if not json_key:
        return f"❌ Invalid statement_type: must be one of {list(STATEMENT_TYPE_KEY_MAP.keys())}"

    url = (
        f"{FINANCIAL_DATASETS_API_BASE}/financials/{statement_type}"
        f"?ticker={ticker}&period={period}&limit={limit}"
        f"&report_period_lte={report_period_lte}&report_period_gte={report_period_gte}"
    )
    response = requests.get(url, headers=headers)
    if response.status_code != 200:
        return f"❌ API Error: {response.status_code} - {response.text}"
    return response.text


# --- Wrapper for Gradio ---
def run_financial_statements(
    ticker, report_period_gte, report_period_lte,
    statement_type, period, limit, api_key
):
    result = get_financial_statements.invoke({
        "ticker": ticker.strip(),
        "report_period_gte": report_period_gte,
        "report_period_lte": report_period_lte,
        "statement_type": statement_type,
        "period": period,
        "limit": int(limit),
        "api_key": api_key.strip()
    })
    return result


# --- Gradio UI ---
with gr.Blocks(theme=gr.themes.Soft(), css=".gr-button { font-size: 16px }") as demo:
    gr.Markdown(
        """
        # 📊 Financial Insights Assistant

        Welcome to the **Financial News & Statements Explorer**!

        🔍 Search for company financials such as income statements, balance sheets, or cash flow data  
        using official ticker symbols (e.g. `AAPL`, `MSFT`, `TSLA`).  
        👉 You’ll need a free API key from [financialdatasets.ai](https://financialdatasets.ai) to use this tool.

        ---
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 🔍 Query Inputs")
            ticker = gr.Textbox(label="Ticker Symbol", placeholder="e.g. AAPL")
            report_gte = gr.Textbox(label="Start Date (YYYY-MM-DD)", value="2020-01-01")
            report_lte = gr.Textbox(label="End Date (YYYY-MM-DD)", value="2024-12-31")
            statement_type = gr.Dropdown(
                ["all", "income-statements", "balance-sheets", "cash-flow-statements"],
                label="Statement Type", value="all"
            )
            period = gr.Radio(["annual", "quarterly", "ttm"], label="Period", value="annual")
            limit = gr.Slider(1, 10, value=4, step=1, label="Limit")
            api_key = gr.Textbox(
                            label="FinancialDatasets API Key",
                            placeholder="Enter your API key (Get one at financialdatasets.ai)",
                            type="password"
                                        )
            gr.Markdown(
                "🔑 Don’t have an API key? [Get one here](https://financialdatasets.ai)"
            )


            submit = gr.Button("📈 Get Financials", variant="primary")

        with gr.Column(scale=1):
            gr.Markdown("### 📄 Output")
            output = gr.Code(label="JSON Result", language="json")

    # Hook
    submit.click(
        fn=run_financial_statements,
        inputs=[ticker, report_gte, report_lte, statement_type, period, limit, api_key],
        outputs=output
    )

    # Footer
    gr.Markdown("---")
    gr.Markdown("Built with ❤️ using [LangChain](https://www.langchain.com) and [FinancialDatasets.ai](https://financialdatasets.ai)")

# Run it
if __name__ == "__main__":
    demo.launch(mcp_server=True)