from typing import Literal

from fastapi import APIRouter, Depends, Request
from pydantic import BaseModel, Field

from private_gpt.chat.extensions.context_filter import ContextFilter
from private_gpt.events.models import ResultContentBlockType
from private_gpt.server.tools.tool_service import ToolService
from private_gpt.server.utils.artifact_input import ArtifactType
from private_gpt.server.utils.auth import authenticated

tool_router = APIRouter(
    prefix="/v1",
    dependencies=[Depends(authenticated)],
    tags=["Tools"],
    responses={401: {"description": "Unauthorized"}},
)


class SemanticSearchBody(BaseModel):
    """Request body for semantic search."""

    query: str = Field(
        ...,
        description="The natural language query to search for relevant content.",
        examples=["What were the Q4 revenue trends?"],
    )
    context_filter: ContextFilter = Field(
        ...,
        description="Filters to narrow the search context (collections, artifacts, metadata).",
    )
    format: Literal["default", "citations"] = Field(
        default="default",
        description=(
            "Format of the result content.\n"
            "'default' returns standard content blocks.\n"
            "'citations' returns blocks annotated for citation formatting."
        ),
        examples=["default"],
    )


class TabularDataAnalysisBody(BaseModel):
    """Request body for tabular data analysis using the tool."""

    query: str = Field(
        ...,
        description="The analysis question or prompt related to tabular data.",
        examples=["Show average revenue by department over the last 3 quarters."],
    )
    context_filter: ContextFilter = Field(
        ...,
        description="Filters to select specific context from ingested documents.",
    )


class DatabaseQueryBody(BaseModel):
    """Request body for database query using the tool."""

    query: str = Field(
        ...,
        description="The natural language query to run against connected databases.",
        examples=["What were the Q4 revenue trends?"],
    )
    artifacts: list[ArtifactType] = Field(
        ...,
        description=(
            "List of SQL database artifacts to query against. "
            "At least one artifact of type 'sql_database' is required."
        ),
    )


class WebSearchQueryBody(BaseModel):
    """Request body for web search using the tool."""

    query: str = Field(
        ...,
        description="The natural language query to search the web.",
        examples=["Latest news on AI advancements"],
    )


class WebFetchBody(BaseModel):
    """Request body for web content fetching using the tool."""

    url: str = Field(
        ...,
        description="The URL of the web page to fetch and extract content from.",
        examples=["https://example.com/article"],
    )


class ToolResponse(BaseModel):
    """Response returned from tool-based operations."""

    content: list[ResultContentBlockType] = Field(
        ...,
        description=(
            "List of content blocks generated by the tool. "
            "Blocks can include plain text, citations, source attributions, or images."
        ),
    )
    is_error: bool = Field(
        False,
        description="True if the tool encountered an error during execution.",
    )


@tool_router.post(
    "/tools/semantic-search",
    response_model=ToolResponse,
    summary="Semantic Search",
    description="Run a semantic search using natural language and contextual filters. Supports citation formatting.",
    responses={
        200: {
            "description": "Semantic search completed successfully.",
            "content": {
                "application/json": {
                    "examples": {
                        "rich_semantic_response": {
                            "summary": "Search result with sources",
                            "value": {
                                "content": [
                                    {
                                        "type": "source",
                                        "sources": [
                                            {
                                                "object": "context.chunk",
                                                "id": "d3ab07fe-b7ab-4fda-8524-1cb41b01200a",
                                                "score": 0.83,
                                                "document": {
                                                    "object": "ingest.document",
                                                    "artifact": "0196ee69-666a-7bcd-800b-b01fa9b2f24c",
                                                    "doc_metadata": {
                                                        "file_name": "dense_x_retrieval.pdf"
                                                    },
                                                },
                                                "text": "Wenhao Yu et al. 2023. Chain-ofnote: Enhancing robustness in retrieval-augmented LMs.",
                                                "content_type": "text/markdown",
                                                "metadata": {
                                                    "page": 13,
                                                    "shorter_id": "GONG",
                                                },
                                                "previous_texts": [],
                                                "next_texts": [],
                                            }
                                        ],
                                    },
                                    {
                                        "type": "text",
                                        "text": "**Context Information**:\nCitation identifier [GONG]\n---\nContent:\nWenhao Yu, Hongming Zhang... Chain-ofnote: Enhancing robustness in retrieval-augmented language models.",
                                    },
                                ],
                                "is_error": False,
                            },
                        },
                    }
                }
            },
        },
        422: {
            "description": "Validation Error - Invalid request parameters",
            "content": {
                "application/json": {
                    "examples": {
                        "missing_query": {
                            "summary": "Missing query field",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "query"],
                                        "msg": "field required",
                                        "type": "value_error.missing",
                                    }
                                ]
                            },
                        },
                        "invalid_format": {
                            "summary": "Invalid format value",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "format"],
                                        "msg": "value is not a valid enumeration member; permitted: 'default', 'citations'",
                                        "type": "type_error.enum",
                                    }
                                ]
                            },
                        },
                    }
                }
            },
        },
    },
    openapi_extra={
        "requestBody": {
            "required": True,
            "description": (
                "Request body for semantic search. Includes the query, optional filters, "
                "and response formatting preferences."
            ),
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/SemanticSearchBody"}
                }
            },
        }
    },
)
async def semantic_search(
    request: Request,
    body: SemanticSearchBody,
) -> ToolResponse:
    """Execute a semantic search using the tool system.

    Performs semantic search against the indexed document corpus using the provided
    query and optional context filters. The format parameter allows formatting
    results for citation-based use cases.

    - Supports natural language search
    - Can return content in citation-ready format
    - Context filters help narrow down relevant data
    - Return ContentBlocks as MCP does
       - It returns Zylon custom content blocks, be careful if you use directly with MCP
    """
    service: ToolService = request.state.injector.get(ToolService)
    result = await service.semantic_search_tool(
        body.query,
        context_filter=body.context_filter,
        generate_citations=body.format == "citations",
    )
    return ToolResponse(content=result.content, is_error=result.is_error)


@tool_router.post(
    "/tools/tabular-data-analysis",
    response_model=ToolResponse,
    summary="Tabular Data Analysis",
    description="Analyze structured tabular data using a natural language query.",
    responses={
        200: {
            "description": "Tabular analysis completed successfully.",
            "content": {
                "application/json": {
                    "examples": {
                        "successful_analysis": {
                            "summary": "Simple text result from tabular analysis",
                            "value": {
                                "content": [
                                    {
                                        "type": "text",
                                        "text": "Average revenue across departments: Sales - $1.2M, Marketing - $850K...",
                                    }
                                ],
                                "is_error": False,
                            },
                        },
                        "text_and_image_output": {
                            "summary": "Text + image result from tabular analysis",
                            "value": {
                                "content": [
                                    {
                                        "type": "text",
                                        "text": "The bar chart shows that Q4 revenue was highest in the Sales department.",
                                    },
                                    {
                                        "type": "image",
                                        "data": "iVBORw0KGgoAAAANSUhEUgAAAAUA"  # Truncated example
                                        "AAAFCAYAAACNbyblAAAAHElEQVQI12P4"
                                        "//8/w38GIAXDIBKE0DHxgljNBAAO"
                                        "9TXL0Y4OHwAAAABJRU5ErkJggg==",
                                        "mime_type": "image/png",
                                    },
                                ],
                                "is_error": False,
                            },
                        },
                    }
                }
            },
        },
        422: {
            "description": "Validation Error - Invalid request parameters",
            "content": {
                "application/json": {
                    "examples": {
                        "missing_query": {
                            "summary": "Missing analysis query",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "query"],
                                        "msg": "field required",
                                        "type": "value_error.missing",
                                    }
                                ]
                            },
                        },
                    }
                }
            },
        },
    },
    openapi_extra={
        "requestBody": {
            "required": True,
            "description": (
                "Request body for analyzing tabular data using a natural language query. "
                "Includes filters to scope the documents containing structured data."
            ),
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/TabularDataAnalysisBody"}
                }
            },
        }
    },
)
async def tabular_data_analysis(
    request: Request,
    body: TabularDataAnalysisBody,
) -> ToolResponse:
    """Run tabular data analysis using the tool system.

    Accepts a question about tabular data and returns structured insights based on
    filtered documents.

    - Ideal for structured datasets as CSVs, spreadsheets, or databases
    - Results may include tables, aggregations, interpreted text, or charts
    - Return ContentBlocks as MCP does
       - It returns Zylon custom content blocks, be careful if you use directly with MCP
    """
    service: ToolService = request.state.injector.get(ToolService)
    result = await service.tabular_data_analysis_tool(
        body.query,
        context_filter=body.context_filter,
    )
    return ToolResponse(content=result.content, is_error=result.is_error)


@tool_router.post(
    "/tools/database-query",
    response_model=ToolResponse,
    summary="Database Query",
    description="Run a natural language query against connected SQL databases.",
    responses={
        200: {
            "description": "Database query completed successfully.",
            "content": {
                "application/json": {
                    "examples": {
                        "successful_query": {
                            "summary": "Simple text result from database query",
                            "value": {
                                "content": [
                                    {
                                        "type": "text",
                                        "text": "The total revenue for Q4 was $3.5M, with the highest contributions from the Sales department.",
                                    }
                                ],
                                "is_error": False,
                            },
                        },
                        "query_with_table_result": {
                            "summary": "Text + table result from database query",
                            "value": {
                                "content": [
                                    {
                                        "type": "text",
                                        "text": "Here is the breakdown of revenue by department for Q4:",
                                    },
                                    {
                                        "type": "table",
                                        "data": {
                                            "columns": ["Department", "Revenue"],
                                            "rows": [
                                                ["Sales", "$1.2M"],
                                                ["Marketing", "$850K"],
                                                ["Engineering", "$1.45M"],
                                            ],
                                        },
                                    },
                                ],
                                "is_error": False,
                            },
                        },
                    }
                }
            },
        },
        422: {
            "description": "Validation Error - Invalid request parameters",
            "content": {
                "application/json": {
                    "examples": {
                        "missing_query": {
                            "summary": "Missing query field",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "query"],
                                        "msg": "field required",
                                        "type": "value_error.missing",
                                    }
                                ]
                            },
                        },
                    }
                }
            },
        },
    },
    openapi_extra={
        "requestBody": {
            "required": True,
            "description": (
                "Request body for querying connected SQL databases using natural language. "
                "Requires at least one SQL database artifact in the tool context."
            ),
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/DatabaseQueryBody"}
                }
            },
        }
    },
)
async def database_query(
    request: Request,
    body: DatabaseQueryBody,
) -> ToolResponse:
    """Run a natural language query against connected SQL databases.

    - Requires at least one SqlDatabaseArtifact in the tool context
    - Returns text summaries, tables, or charts based on query results
    - Return ContentBlocks as MCP does
       - It returns Zylon custom content blocks, be careful if you use directly with MCP
    """
    service: ToolService = request.state.injector.get(ToolService)
    result = await service.database_query_tool(
        body.query,
        sql_artifacts=[
            ctx
            for ctx in body.artifacts
            if isinstance(ctx, ArtifactType) and ctx.type == "sql_database"
        ],
    )
    return ToolResponse(content=result.content, is_error=result.is_error)


@tool_router.post(
    "/tools/web-fetch",
    response_model=ToolResponse,
    summary="Web Fetch",
    description="Fetch and extract content from a specified web URL.",
    responses={
        200: {
            "description": "Web content fetched successfully.",
            "content": {
                "application/json": {
                    "examples": {
                        "successful_fetch": {
                            "summary": "Fetched text content from a web page",
                            "value": {
                                "content": [
                                    {
                                        "type": "text",
                                        "text": "This is the main content extracted from the web page...",
                                    }
                                ],
                                "is_error": False,
                            },
                        },
                    }
                }
            },
        },
        422: {
            "description": "Validation Error - Invalid request parameters",
            "content": {
                "application/json": {
                    "examples": {
                        "missing_url": {
                            "summary": "Missing URL in the last user message",
                            "value": {
                                "detail": [
                                    {
                                        "loc": ["body", "context_filter"],
                                        "msg": "No URL found in the last user message.",
                                        "type": "value_error",
                                    }
                                ]
                            },
                        },
                    }
                }
            },
        },
    },
    openapi_extra={
        "requestBody": {
            "required": True,
            "description": (
                "Request body for fetching web content. The last user message must contain a valid URL."
            ),
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/WebFetchBody"}
                }
            },
        }
    },
)
async def web_fetch(
    request: Request,
    body: WebFetchBody,
) -> ToolResponse:
    """Fetch and extract content from a specified web URL.

    - Returns the main textual content of the web page in Markdown
    - Return ContentBlocks as MCP does
       - It returns Zylon custom content blocks, be careful if you use directly with MCP
    """
    service: ToolService = request.state.injector.get(ToolService)
    result = await service.web_fetch_tool(body.url)
    return ToolResponse(content=result.content, is_error=result.is_error)


@tool_router.post(
    "/tools/web-search",
    response_model=ToolResponse,
    summary="Web Search",
    description="Search the web for information related to a natural language query and return aggregated results from multiple sources.",
    responses={
        200: {
            "description": "Web search completed successfully.",
            "content": {
                "application/json": {
                    "examples": {
                        "successful_search": {
                            "summary": "Web search results with sources and text content",
                            "value": {
                                "content": [
                                    {
                                        "type": "source",
                                        "sources": [
                                            {
                                                "id": "website_a1f2ae32-e661-4f61-ad94-7f26163b51da",
                                                "object": "context.website",
                                                "url": "https://example.com/page",
                                                "favicon_url": "https://example.com/favicon.png",
                                                "title": "Example Page Title",
                                                "description": "Brief description of the page content",
                                                "metadata": None,
                                                "content_type": "text/markdown",
                                                "content": "Full markdown content from the source...",
                                            }
                                        ],
                                    },
                                    {
                                        "type": "text",
                                        "text": "1. Example Page Title\nDescription: Brief description...\nURL: https://example.com/page\nAge: Unknown\nContent: Full content preview...",
                                    },
                                ],
                                "is_error": False,
                            },
                        },
                    }
                }
            },
        },
        422: {
            "description": "Validation Error - Invalid request parameters",
            "content": {
                "application/json": {
                    "examples": {
                        "missing_query": {
                            "summary": "Missing required query field",
                            "value": {
                                "detail": [
                                    {
                                        "type": "missing",
                                        "loc": ["body", "query"],
                                        "msg": "Field required",
                                        "input": {},
                                    }
                                ]
                            },
                        },
                    }
                }
            },
        },
    },
    openapi_extra={
        "requestBody": {
            "required": True,
            "description": (
                "Request body for web search. Includes the natural language query to search the web."
            ),
            "content": {
                "application/json": {
                    "schema": {"$ref": "#/components/schemas/WebSearchQueryBody"}
                }
            },
        }
    },
)
async def web_search(
    request: Request,
    body: WebSearchQueryBody,
) -> ToolResponse:
    """Search the web and get information from multiple sources based on the query.

    Performs a web search using the provided natural language query and returns
    aggregated results from multiple web sources including:
    - Source metadata (URL, title, description, favicon)
    - Full markdown content from each source
    - Formatted text summaries with content previews

    Returns search results in a structured format with both source objects
    and text representations for easy consumption.
    """
    service: ToolService = request.state.injector.get(ToolService)
    result = await service.web_search_tool(body.query)
    return ToolResponse(content=result.content, is_error=result.is_error)
