"""Spice Bae - AI-Powered Spice Advisor with MCP Server.

This application provides:
1. Web UI for querying spice information
2. MCP server at /gradio_api/mcp/ for AI agent integration
3. Conversational AI chat interface (LlamaIndex + Claude)

All data is sourced from USDA FoodData Central with full attribution.
"""

import os
import traceback
from functools import wraps
from typing import List, Tuple, Callable, Any
from dotenv import load_dotenv
import gradio as gr


def handle_errors(func: Callable) -> Callable:
    """Decorator to handle errors gracefully in UI functions.

    Args:
        func: Function to wrap with error handling.

    Returns:
        Wrapped function that catches exceptions.
    """
    @wraps(func)
    def wrapper(*args, **kwargs) -> Any:
        try:
            return func(*args, **kwargs)
        except ConnectionError:
            return None, "Connection error. Please check your internet connection and try again."
        except TimeoutError:
            return None, "Request timed out. The server may be busy. Please try again."
        except Exception as e:
            error_msg = f"An error occurred: {str(e)}"
            print(f"[ERROR] {func.__name__}: {traceback.format_exc()}")
            if hasattr(func, '__annotations__'):
                return_type = func.__annotations__.get('return', None)
                if return_type and 'Tuple' in str(return_type):
                    return None, error_msg
            return error_msg
    return wrapper


from tools.mcp_tools import (
    get_spice_information as _get_spice_information,
    get_spice_info_with_image as _get_spice_info_with_image,
    get_health_benefits_with_image as _get_health_benefits_with_image,
    get_safety_info_with_image as _get_safety_info_with_image,
    list_available_spices as _list_available_spices,
    get_nutrient_content as _get_nutrient_content,
    compare_spices as _compare_spices,
    find_spice_substitutes as _find_spice_substitutes,
    get_health_benefits_info as _get_health_benefits_info,
    find_spices_for_health_benefit as _find_spices_for_health_benefit,
    get_spice_safety_information as _get_spice_safety_information,
    find_medicinal_substitutes as _find_medicinal_substitutes
)
from tools.llama_agent import SpiceAgent


# Wrap all functions with error handling
@handle_errors
def get_spice_information(spice_name: str) -> str:
    return _get_spice_information(spice_name)


@handle_errors
def get_spice_info_with_image(spice_name: str) -> Tuple[str, str]:
    return _get_spice_info_with_image(spice_name)


@handle_errors
def get_health_benefits_with_image(spice_name: str) -> Tuple[str, str]:
    return _get_health_benefits_with_image(spice_name)


@handle_errors
def get_safety_info_with_image(spice_name: str) -> Tuple[str, str]:
    return _get_safety_info_with_image(spice_name)


@handle_errors
def list_available_spices() -> str:
    return _list_available_spices()


@handle_errors
def get_nutrient_content(spice_name: str, nutrient_name: str) -> str:
    return _get_nutrient_content(spice_name, nutrient_name)


@handle_errors
def compare_spices(spice1: str, spice2: str) -> str:
    return _compare_spices(spice1, spice2)


@handle_errors
def find_spice_substitutes(spice_name: str, limit: int = 5) -> str:
    return _find_spice_substitutes(spice_name, limit)


@handle_errors
def get_health_benefits_info(spice_name: str) -> str:
    return _get_health_benefits_info(spice_name)


@handle_errors
def find_spices_for_health_benefit(benefit: str) -> str:
    return _find_spices_for_health_benefit(benefit)


@handle_errors
def get_spice_safety_information(spice_name: str) -> str:
    return _get_spice_safety_information(spice_name)


if not os.getenv("SPACE_ID"):
    load_dotenv()

# Initialize LlamaIndex agent (lazy initialization)
spice_agent = None


def get_agent() -> SpiceAgent:
    """Get or create the spice agent instance."""
    global spice_agent
    if spice_agent is None:
        spice_agent = SpiceAgent()
    return spice_agent


def get_usage_stats() -> str:
    """Get usage statistics from the guardrails.

    Returns:
        Formatted usage statistics string.
    """
    agent = get_agent()

    if not agent.guardrails:
        return "Guardrails are disabled. No usage tracking available."

    usage_guardrail = agent.guardrails.get_guardrail("usage_tracking")

    if not usage_guardrail:
        return "Usage tracking guardrail not found."

    summary = usage_guardrail.get_daily_summary()

    output = []
    output.append("=== USAGE DASHBOARD ===\n")
    output.append(f"Date: {summary['date']}")
    output.append(f"Total Requests: {summary['total_requests']}")
    output.append(f"\n--- Token Usage ---")
    output.append(f"Input Tokens: {summary['total_input_tokens']:,}")
    output.append(f"Output Tokens: {summary['total_output_tokens']:,}")
    output.append(f"Total Tokens: {summary['total_input_tokens'] + summary['total_output_tokens']:,}")
    output.append(f"\n--- Cost ---")
    output.append(f"Today's Cost: ${summary['total_cost_usd']:.4f}")
    output.append(f"Daily Limit: ${summary['daily_limit_usd']:.2f}")
    output.append(f"Remaining Budget: ${summary['remaining_budget_usd']:.4f}")

    pct_used = (summary['total_cost_usd'] / summary['daily_limit_usd']) * 100 if summary['daily_limit_usd'] > 0 else 0
    output.append(f"Budget Used: {pct_used:.1f}%")

    return "\n".join(output)


def chat_response(
    message: str,
    history: List[Tuple[str, str]]
) -> Tuple[str, List[Tuple[str, str]]]:
    """Process chat message and return response.

    Args:
        message: User's input message.
        history: List of (user, assistant) message tuples.

    Returns:
        Tuple of (empty string for input, updated history).
    """
    agent = get_agent()

    if not agent.is_ready():
        response = (
            "The AI chat feature requires an Anthropic API key (ANTHROPIC_API_KEY). "
            "Please use the individual tabs above for spice queries, or set "
            "the ANTHROPIC_API_KEY environment variable to enable conversational AI."
        )
    else:
        response = agent.chat(message)

    history.append((message, response))
    return "", history


def create_chat_tab():
    """Create the conversational AI chat tab."""
    with gr.Tab("AI Chat"):
        gr.Markdown(
            """
            ### Ask me anything about spices!

            I can help you with:
            - Spice information and nutritional data
            - Health benefits and traditional uses
            - Finding substitutes for spices
            - Safety information and cautions

            **Note:** Requires ANTHROPIC_API_KEY for AI responses. Use other tabs if unavailable.
            """
        )

        chatbot = gr.Chatbot(
            label="Conversation",
            height=400,
            type="tuples",
            show_copy_button=True
        )

        with gr.Row():
            msg_input = gr.Textbox(
                label="Your Question",
                placeholder="e.g., What spices help with inflammation?",
                scale=4,
                show_label=False
            )
            send_btn = gr.Button("Send", variant="primary", scale=1)

        status = gr.Markdown("", elem_classes=["status-ready"])

        gr.Examples(
            examples=[
                ["What are the health benefits of turmeric?"],
                ["Which spices help lower cholesterol?"],
                ["What's a good substitute for cinnamon?"],
                ["Is garlic safe to take with blood thinners?"],
                ["What spices are high in iron?"],
            ],
            inputs=msg_input
        )

        def chat_with_status(message, history):
            """Chat response with status updates."""
            result = chat_response(message, history)
            return result[0], result[1], ""

        msg_input.submit(
            fn=lambda: "Thinking...",
            outputs=status
        ).then(
            fn=chat_with_status,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot, status],
            show_progress="minimal"
        )
        send_btn.click(
            fn=lambda: "Thinking...",
            outputs=status
        ).then(
            fn=chat_with_status,
            inputs=[msg_input, chatbot],
            outputs=[msg_input, chatbot, status],
            show_progress="minimal"
        )


def create_spice_database_tab():
    """Create the consolidated spice database tab."""
    with gr.Tab("Spice Database"):
        gr.Markdown("### Explore spice information, nutrients, and comparisons")

        with gr.Accordion("Spice Lookup", open=True):
            with gr.Row():
                spice_input = gr.Textbox(
                    label="Spice Name",
                    placeholder="e.g., turmeric, ginger, cinnamon",
                    scale=3
                )
                search_btn = gr.Button("Search", variant="primary", scale=1)

            with gr.Row():
                spice_image = gr.Image(
                    label="Spice Image",
                    height=200,
                    width=200,
                    show_label=False,
                    scale=1
                )
                info_output = gr.Textbox(
                    label="Spice Information",
                    lines=10,
                    interactive=False,
                    scale=3
                )

            gr.Examples(
                examples=[["turmeric"], ["ginger root"], ["cinnamon"], ["garlic"]],
                inputs=spice_input
            )

            search_btn.click(
                fn=get_spice_info_with_image,
                inputs=spice_input,
                outputs=[spice_image, info_output],
                show_progress="minimal"
            )
            spice_input.submit(
                fn=get_spice_info_with_image,
                inputs=spice_input,
                outputs=[spice_image, info_output],
                show_progress="minimal"
            )

        with gr.Accordion("Find Substitutes", open=False):
            with gr.Row():
                sub_spice_input = gr.Textbox(
                    label="Spice Name",
                    placeholder="e.g., cinnamon",
                    scale=3
                )
                limit_slider = gr.Slider(
                    minimum=3, maximum=10, value=5, step=1,
                    label="Results", scale=1
                )
                find_btn = gr.Button("Find", variant="primary", scale=1)

            sub_output = gr.Textbox(
                label="Substitute Suggestions",
                lines=12,
                interactive=False
            )

            find_btn.click(
                fn=find_spice_substitutes,
                inputs=[sub_spice_input, limit_slider],
                outputs=sub_output,
                show_progress="minimal"
            )

        with gr.Accordion("Compare Spices", open=False):
            with gr.Row():
                spice1_input = gr.Textbox(label="Spice 1", placeholder="e.g., turmeric", scale=2)
                spice2_input = gr.Textbox(label="Spice 2", placeholder="e.g., ginger", scale=2)
                compare_btn = gr.Button("Compare", variant="primary", scale=1)

            compare_output = gr.Textbox(
                label="Comparison Results",
                lines=15,
                interactive=False
            )

            compare_btn.click(
                fn=compare_spices,
                inputs=[spice1_input, spice2_input],
                outputs=compare_output,
                show_progress="minimal"
            )

        with gr.Accordion("Nutrient Lookup", open=False):
            with gr.Row():
                nut_spice = gr.Textbox(label="Spice", placeholder="e.g., turmeric", scale=2)
                nut_name = gr.Textbox(label="Nutrient", placeholder="e.g., iron", scale=2)
                nut_btn = gr.Button("Search", variant="primary", scale=1)

            nut_output = gr.Textbox(label="Nutrient Information", lines=6, interactive=False)

            nut_btn.click(
                fn=get_nutrient_content,
                inputs=[nut_spice, nut_name],
                outputs=nut_output,
                show_progress="minimal"
            )

        with gr.Accordion("All Available Spices", open=False):
            list_btn = gr.Button("Show All Spices", variant="secondary")
            list_output = gr.Textbox(label="Spice List", lines=12, interactive=False)

            list_btn.click(
                fn=list_available_spices,
                outputs=list_output,
                show_progress="minimal"
            )


def create_health_safety_tab():
    """Create the consolidated health and safety tab."""
    with gr.Tab("Health & Safety"):
        gr.Markdown(
            """
            ### Health benefits and safety information
            Data sourced from **NCCIH** (National Center for Complementary and Integrative Health)
            """
        )

        with gr.Accordion("Health Benefits Lookup", open=True):
            with gr.Row():
                health_spice = gr.Textbox(
                    label="Spice Name",
                    placeholder="e.g., Garlic, Cinnamon",
                    scale=3
                )
                health_btn = gr.Button("Get Health Info", variant="primary", scale=1)

            with gr.Row():
                health_image = gr.Image(
                    label="Spice Image",
                    height=200,
                    width=200,
                    show_label=False,
                    scale=1
                )
                health_output = gr.Textbox(
                    label="Health Information",
                    lines=12,
                    interactive=False,
                    scale=3
                )

            gr.Examples(
                examples=[["Garlic"], ["Cinnamon"], ["Fenugreek"], ["Turmeric"]],
                inputs=health_spice
            )

            health_btn.click(
                fn=get_health_benefits_with_image,
                inputs=health_spice,
                outputs=[health_image, health_output],
                show_progress="minimal"
            )
            health_spice.submit(
                fn=get_health_benefits_with_image,
                inputs=health_spice,
                outputs=[health_image, health_output],
                show_progress="minimal"
            )

        with gr.Accordion("Search by Health Condition", open=False):
            with gr.Row():
                condition_input = gr.Textbox(
                    label="Health Condition",
                    placeholder="e.g., diabetes, cholesterol, inflammation",
                    scale=3
                )
                condition_btn = gr.Button("Search", variant="primary", scale=1)

            condition_output = gr.Textbox(
                label="Matching Spices",
                lines=12,
                interactive=False
            )

            gr.Examples(
                examples=[["diabetes"], ["cholesterol"], ["inflammation"]],
                inputs=condition_input
            )

            condition_btn.click(
                fn=find_spices_for_health_benefit,
                inputs=condition_input,
                outputs=condition_output,
                show_progress="minimal"
            )

        with gr.Accordion("Safety Information", open=False):
            with gr.Row():
                safety_spice = gr.Textbox(
                    label="Spice Name",
                    placeholder="e.g., Garlic, Cinnamon",
                    scale=3
                )
                safety_btn = gr.Button("Get Safety Info", variant="primary", scale=1)

            with gr.Row():
                safety_image = gr.Image(
                    label="Spice Image",
                    height=200,
                    width=200,
                    show_label=False,
                    scale=1
                )
                safety_output = gr.Textbox(
                    label="Safety Information",
                    lines=12,
                    interactive=False,
                    scale=3
                )

            gr.Markdown(
                "**IMPORTANT:** Always consult healthcare providers before using herbs medicinally."
            )

            safety_btn.click(
                fn=get_safety_info_with_image,
                inputs=safety_spice,
                outputs=[safety_image, safety_output],
                show_progress="minimal"
            )
            safety_spice.submit(
                fn=get_safety_info_with_image,
                inputs=safety_spice,
                outputs=[safety_image, safety_output],
                show_progress="minimal"
            )


def create_usage_dashboard_tab():
    """Create the usage dashboard tab."""
    with gr.Tab("Usage"):
        gr.Markdown("### API Usage Dashboard")

        refresh_btn = gr.Button("Refresh Stats", variant="primary")

        stats_output = gr.Textbox(
            label="Usage Statistics",
            lines=12,
            interactive=False,
            value="Click 'Refresh Stats' to view current usage."
        )

        gr.Markdown("*Usage tracked for AI Chat. Stats reset daily.*")

        refresh_btn.click(
            fn=get_usage_stats,
            outputs=stats_output,
            show_progress="minimal"
        )


# Custom CSS for better loading indicators and styling (supports dark mode)
CUSTOM_CSS = """
/* Light mode styles */
.loading-indicator {
    display: flex;
    align-items: center;
    gap: 8px;
    padding: 8px 12px;
    background: linear-gradient(90deg, #f0f7ff 0%, #e8f4fd 100%);
    border-radius: 6px;
    margin: 8px 0;
}
.loading-spinner {
    width: 16px;
    height: 16px;
    border: 2px solid #3b82f6;
    border-top-color: transparent;
    border-radius: 50%;
    animation: spin 0.8s linear infinite;
}
@keyframes spin {
    to { transform: rotate(360deg); }
}
.status-ready { color: #059669; }
.status-loading { color: #3b82f6; }
.status-error { color: #dc2626; }
footer { display: none !important; }

/* Dark mode styles */
@media (prefers-color-scheme: dark) {
    .loading-indicator {
        background: linear-gradient(90deg, #1e3a5f 0%, #1a365d 100%);
    }
    .status-ready { color: #34d399; }
    .status-loading { color: #60a5fa; }
    .status-error { color: #f87171; }
}

/* Gradio dark mode class overrides */
.dark .loading-indicator {
    background: linear-gradient(90deg, #1e3a5f 0%, #1a365d 100%);
}
.dark .status-ready { color: #34d399; }
.dark .status-loading { color: #60a5fa; }
.dark .status-error { color: #f87171; }

/* Improved image styling */
.image-container img {
    border-radius: 8px;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1);
}
.dark .image-container img {
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.3);
}
"""

# Create main application
with gr.Blocks(
    title="Spice Bae - AI Spice Advisor",
    theme=gr.themes.Soft(),
    css=CUSTOM_CSS
) as demo:

    gr.Markdown(
        """
        # Spice Bae

        Your AI-powered spice advisor for nutritional content and health properties.
        All data sourced from **USDA FoodData Central** with full attribution.

        **Note:** This is for educational purposes only, not medical advice.
        """
    )

    # Create consolidated tabs (4 tabs instead of 10)
    create_chat_tab()
    create_spice_database_tab()
    create_health_safety_tab()
    create_usage_dashboard_tab()

    gr.Markdown(
        """
        ---
        **Data Sources:** USDA FoodData Central | NCCIH (Public Domain)

        **Images:** Wikipedia/Wikimedia Commons (CC/Public Domain)

        **MCP Server:** Available at `/gradio_api/mcp/`
        """
    )


if __name__ == "__main__":
    # Launch with MCP server enabled
    demo.launch(
        mcp_server=True,
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False
    )

    print("\n" + "="*60)
    print("Spice Bae is running!")
    print("Web UI: http://localhost:7860")
    print("MCP Server: http://localhost:7860/gradio_api/mcp/sse")
    print("="*60)