from solace_agent_mesh.agent.tools.dynamic_tool import DynamicTool
from google.adk.tools import ToolContext
from pydantic import BaseModel, Field
from typing import Optional
import httpx

class ApiToolConfig(BaseModel):
    """Configuration for the API tool."""
    api_key: str = Field(description="API key for authentication")
    base_url: str = Field(description="Base URL for the API")
    timeout: int = Field(default=30, description="Request timeout in seconds")

class ApiSearchTool(DynamicTool):
    """Tool for searching an external API."""
    
    config_model = ApiToolConfig
    
    def __init__(self, tool_config: Optional[ApiToolConfig] = None):
        super().__init__(tool_config)
        self.client = None
    
    async def init(self, component, tool_config: ApiToolConfig) -> None:
        """Initialize the HTTP client."""
        self.client = httpx.AsyncClient(
            base_url=tool_config.base_url,
            timeout=tool_config.timeout,
            headers={"Authorization": f"Bearer {tool_config.api_key}"}
        )
    
    async def cleanup(self, component, tool_config: ApiToolConfig) -> None:
        """Clean up the HTTP client."""
        if self.client:
            await self.client.aclose()
    
    @property
    def tool_name(self) -> str:
        return "search_api"
    
    @property
    def tool_description(self) -> str:
        return "Search the external API for information."
    
    @property
    def tool_parameters_schema(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query"
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of results",
                    "default": 10
                }
            },
            "required": ["query"]
        }
    
    async def execute(
        self,
        query: str,
        limit: int = 10,
        tool_context: ToolContext = None
    ) -> dict:
        """Execute the API search."""
        if not self.client:
            return {"error": "Tool not initialized"}
        
        try:
            response = await self.client.get(
                "/search",
                params={"q": query, "limit": limit}
            )
            response.raise_for_status()
            
            return {
                "status": "success",
                "results": response.json(),
                "count": len(response.json())
            }
        
        except Exception as e:
            return {
                "status": "error",
                "message": f"Search failed: {str(e)}"
            }

from typing import Optional, List, Dict, Any
from google.adk.tools import ToolContext
from google.genai import types as adk_types
import sqlite3
import json
import pandas as pd

def execute_sql_query(
    query: str,
    tool_context: ToolContext = None
) -> dict:
    """
    Execute a SQL query against the database.
    
    This tool connects to the configured SQLite database and executes
    the provided SQL query. For large result sets (>100 rows), the
    results are saved as a CSV artifact.
    
    Args:
        query: The SQL query to execute
        tool_context: The ADK tool context
    
    Returns:
        Dictionary with query results or error message
    """
    if not tool_context:
        return {
            "status": "error",
            "message": "ToolContext is required for this tool"
        }
    
    # Get database path from component config
    inv_context = tool_context._invocation_context
    component = getattr(inv_context, 'component', None)
    
    if not component or not hasattr(component, 'db_connection'):
        return {
            "status": "error",
            "message": "Database not initialized"
        }
    
    try:
        # Execute query
        conn = component.db_connection
        df = pd.read_sql_query(query, conn)
        
        row_count = len(df)
        
        # For small results, return directly
        if row_count <= 100:
            return {
                "status": "success",
                "row_count": row_count,
                "columns": df.columns.tolist(),
                "data": df.to_dict(orient='records')
            }
        
        # For large results, save as artifact
        csv_content = df.to_csv(index=False)
        artifact_part = adk_types.Part.from_text(
            data=csv_content,
            mime_type="text/csv"
        )
        
        filename = "query_results.csv"
        version = tool_context.save_artifact(
            filename=filename,
            artifact=artifact_part
        )
        
        return {
            "status": "success",
            "row_count": row_count,
            "columns": df.columns.tolist(),
            "artifact": filename,
            "version": version,
            "message": f"Query returned {row_count} rows. Results saved to {filename}",
            "preview": df.head(5).to_dict(orient='records')
        }
    
    except Exception as e:
        return {
            "status": "error",
            "message": f"Query failed: {str(e)}",
            "query": query
        }