from typing import Any, Dict, Tuple
from functools import partial

from mcp import types as mcp_types

from ..logging import get_logger
from ..meta_client import MetaClient
from ..schemas import (
    BudgetAllocationOptimizationRequest,
    BidOptimizationRequest,
    CampaignMixOptimizationRequest,
    DaypartingOptimizationRequest,
    AudiencePrioritizationRequest,
    ToolResult
)
from .base import ToolHandler, _tool_schema

logger = get_logger(__name__)

async def optimize_budget_allocation_handler(client: MetaClient, arguments: Dict[str, Any]) -> ToolResult:
    request = BudgetAllocationOptimizationRequest.model_validate(arguments)
    try:
        from ..optimization_algorithms import get_optimization_engine
        engine = get_optimization_engine()
        
        data = await engine.optimize_budget_allocation(request)
        
        return ToolResult(ok=True, data=data)
    except Exception as exc:
        logger.error("tool_optimize_budget_allocation_failed", error=str(exc))
        return ToolResult(ok=False, error={"message": str(exc)})

async def optimize_bids_handler(client: MetaClient, arguments: Dict[str, Any]) -> ToolResult:
    request = BidOptimizationRequest.model_validate(arguments)
    try:
        from ..optimization_algorithms import get_optimization_engine
        engine = get_optimization_engine()
        
        data = await engine.optimize_bids(request)
        
        return ToolResult(ok=True, data=data)
    except Exception as exc:
        logger.error("tool_optimize_bids_failed", error=str(exc))
        return ToolResult(ok=False, error={"message": str(exc)})

async def optimize_campaign_mix_handler(client: MetaClient, arguments: Dict[str, Any]) -> ToolResult:
    request = CampaignMixOptimizationRequest.model_validate(arguments)
    try:
        from ..optimization_algorithms import get_optimization_engine
        engine = get_optimization_engine()
        
        data = await engine.optimize_campaign_mix(request)
        
        return ToolResult(ok=True, data=data)
    except Exception as exc:
        logger.error("tool_optimize_campaign_mix_failed", error=str(exc))
        return ToolResult(ok=False, error={"message": str(exc)})

async def optimize_dayparting_handler(client: MetaClient, arguments: Dict[str, Any]) -> ToolResult:
    request = DaypartingOptimizationRequest.model_validate(arguments)
    try:
        from ..optimization_algorithms import get_optimization_engine
        engine = get_optimization_engine()
        
        data = await engine.optimize_dayparting(request)
        
        return ToolResult(ok=True, data=data)
    except Exception as exc:
        logger.error("tool_optimize_dayparting_failed", error=str(exc))
        return ToolResult(ok=False, error={"message": str(exc)})

async def optimize_audience_prioritization_handler(client: MetaClient, arguments: Dict[str, Any]) -> ToolResult:
    request = AudiencePrioritizationRequest.model_validate(arguments)
    try:
        from ..optimization_algorithms import get_optimization_engine
        engine = get_optimization_engine()
        
        data = await engine.prioritize_audiences(request)
        
        return ToolResult(ok=True, data=data)
    except Exception as exc:
        logger.error("tool_optimize_audience_prioritization_failed", error=str(exc))
        return ToolResult(ok=False, error={"message": str(exc)})

def get_tools(client: MetaClient) -> Dict[str, Tuple[mcp_types.Tool, ToolHandler]]:
    return {
        "meta.optimize.budget_allocation": (
            mcp_types.Tool(
                name="meta.optimize.budget_allocation",
                description="Optimize budget allocation across campaigns using ML algorithms. Reallocate spend from low to high performers. Use for automated portfolio budget management.",
                inputSchema=_tool_schema(BudgetAllocationOptimizationRequest),
            ),
            partial(optimize_budget_allocation_handler, client),
        ),
        "meta.optimize.bids": (
            mcp_types.Tool(
                name="meta.optimize.bids",
                description="Optimize bid strategy using historical performance data. Recommends bid adjustments by audience, placement, and time. Use for algorithmic bid management.",
                inputSchema=_tool_schema(BidOptimizationRequest),
            ),
            partial(optimize_bids_handler, client),
        ),
        "meta.optimize.campaign_mix": (
            mcp_types.Tool(
                name="meta.optimize.campaign_mix",
                description="Optimize campaign mix across objectives (awareness, traffic, conversions). Balance funnel stages for optimal overall performance. Use for full-funnel strategy.",
                inputSchema=_tool_schema(CampaignMixOptimizationRequest),
            ),
            partial(optimize_campaign_mix_handler, client),
        ),
        "meta.optimize.dayparting": (
            mcp_types.Tool(
                name="meta.optimize.dayparting",
                description="Optimize dayparting schedule based on hourly performance patterns. Find best hours and days to run ads. Use for time-based bidding and scheduling.",
                inputSchema=_tool_schema(DaypartingOptimizationRequest),
            ),
            partial(optimize_dayparting_handler, client),
        ),
        "meta.optimize.audience_prioritization": (
            mcp_types.Tool(
                name="meta.optimize.audience_prioritization",
                description="Prioritize audiences by predicted conversion value. Rank custom and lookalike audiences for budget focus. Use for audience-based optimization.",
                inputSchema=_tool_schema(AudiencePrioritizationRequest),
            ),
            partial(optimize_audience_prioritization_handler, client),
        ),
    }
