from fastapi import APIRouter

from axon.models import (
    CrawlRequest,
    ExtractRequest,
    MapRequest,
    SearchRequest,
    TavilyExtractRequest,
    TavilySearchRequest,
)
from axon.routers.crawl import crawl
from axon.routers.extract import extract
from axon.routers.map import map_site
from axon.routers.search import search

router = APIRouter(prefix="/tavily")


@router.post("/search")
async def tavily_search(request: TavilySearchRequest):
    include_answer = request.include_answer not in (False, "false", "none", "basic")
    response = await search(
        SearchRequest(
            query=request.query,
            max_results=request.max_results,
            search_depth=request.search_depth,
            include_answer=include_answer,
            include_raw_content=request.include_raw_content,
            domains_whitelist=request.include_domains,
            domains_blacklist=request.exclude_domains,
            topic=request.topic,
            time_range=request.time_range,
        )
    )
    return response.model_dump()


@router.post("/extract")
async def tavily_extract(request: TavilyExtractRequest):
    urls = [request.urls] if isinstance(request.urls, str) else request.urls
    response = await extract(ExtractRequest(urls=urls, include_images=request.include_images))
    return response.model_dump()


@router.post("/crawl")
async def tavily_crawl(request: CrawlRequest):
    return (await crawl(request)).model_dump()


@router.post("/map")
async def tavily_map(request: MapRequest):
    return (await map_site(request)).model_dump()