import requests
import json
import argparse
import sys
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich import box

# Initialize rich console
console = Console()

def query_api(query="rainfall data north east india", api_url="http://localhost:8000"):
    """
    Query the API and return the response
    """
    try:
        # First check API health
        health_response = requests.get(f"{api_url}/health")
        health_data = health_response.json()
        
        if health_response.status_code != 200 or health_data.get("status") != "ok":
            console.print(Panel(f"[bold red]API is not healthy: {health_data}", 
                               title="Error", border_style="red"))
            return None
            
        # Query the API
        payload = {
            "query": query,
            "session_id": "view-tool-session"
        }
        
        response = requests.post(f"{api_url}/api/query", json=payload)
        if response.status_code != 200:
            console.print(Panel(f"[bold red]API request failed with status code {response.status_code}", 
                               title="Error", border_style="red"))
            return None
            
        return response.json()
        
    except requests.exceptions.ConnectionError:
        console.print(Panel(f"[bold red]Failed to connect to API at {api_url}. Make sure the backend server is running.", 
                           title="Connection Error", border_style="red"))
        return None
    except Exception as e:
        console.print(Panel(f"[bold red]Error querying API: {str(e)}", 
                           title="Error", border_style="red"))
        return None

def get_rainfall_data(region="north east india", api_url="http://localhost:8000"):
    """
    Get rainfall data for a specific region
    """
    query = f"rainfall data {region}"
    return query_api(query, api_url)
    
def display_api_response(response):
    """
    Display the API response in a formatted way
    """
    if not response:
        return
        
    console.print(Panel(response.get("answer", "No answer provided"), 
                       title="AI Assistant Answer", 
                       border_style="blue"))
    
    sources = response.get("sources", [])
    if sources:
        source_table = Table(title="Sources", box=box.ROUNDED)
        source_table.add_column("Source URL", style="cyan")
        
        for source in sources:
            source_table.add_row(source)
            
        console.print(source_table)
        
def display_visualization_data(region="north east india", parameter="Rainfall", api_url="http://localhost:8000"):
    """
    Get and display visualization data for a specific region and parameter
    """
    try:
        payload = {
            "parameter": parameter,
            "satellite": "INSAT-3D",
            "region": region,
            "dateFrom": None,
            "dateTo": None
        }
        
        response = requests.post(f"{api_url}/api/visualization", json=payload)
        if response.status_code != 200:
            console.print(Panel(f"[bold red]Visualization API request failed with status code {response.status_code}", 
                               title="Error", border_style="red"))
            return
            
        data = response.json()
        
        # Display the rainfall data in a table
        table = Table(title=f"{parameter} Data for {region}", box=box.ROUNDED)
        table.add_column("Date", style="cyan")
        table.add_column("Value (mm)", style="green")
        
        for point in data.get("data", []):
            table.add_row(point.get("date", "Unknown"), str(point.get("value", "N/A")))
            
        console.print(table)
        
    except Exception as e:
        console.print(Panel(f"[bold red]Error retrieving visualization data: {str(e)}", 
                           title="Error", border_style="red"))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="View data from the MOSDAC Vector Store via API")
    
    parser.add_argument("--api-url", type=str, default="http://localhost:8000", 
                        help="URL of the backend API")
    parser.add_argument("--query", type=str, default="rainfall data north east india", 
                        help="Query to search for in the vector store")
    parser.add_argument("--region", type=str, default="North East India", 
                        help="Region to get data for")
    parser.add_argument("--parameter", type=str, default="Rainfall", 
                        help="Parameter to get data for (Rainfall, Sea Surface Temperature, etc.)")
    parser.add_argument("--visualization", action="store_true", 
                        help="Get visualization data instead of querying the AI assistant")
    
    args = parser.parse_args()
    
    console.print(Panel(f"[bold green]MOSDAC Data Viewer", 
                       subtitle="Access vector store data via API", 
                       border_style="green"))
    
    if args.visualization:
        console.print(f"[yellow]Getting visualization data for {args.parameter} in {args.region}...")
        display_visualization_data(args.region, args.parameter, args.api_url)
    else:
        console.print(f"[yellow]Querying API with: [bold]{args.query}[/bold]")
        response = query_api(args.query, args.api_url)
        display_api_response(response)