"""
Client implementation for querying metrics from the monitoring system.
"""
import json
import csv
import requests
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, Any, List, Optional

# Import logger_utils
from app.logcli.utils.logger_utils import log_info, log_error

# Import auth for getting credentials
from auth.client_auth import get_credentials


def get_metrics(
    cluster_id: str,
    start_time: datetime,
    end_time: datetime,
    metrics: List[str],
    instance: Optional[str] = None,
    step: str = "15s",
    region: Optional[str] = None,
    o11y_id: Optional[str] = None,
    tenant_id: Optional[str] = None,
    vendor: str = "aws"
) -> Dict[str, Any]:
    """
    Query metrics from the observability endpoint.
    
    Args:
        cluster_id: The cluster ID to query metrics for
        start_time: Start time for the query
        end_time: End time for the query
        metrics: List of metric names to query
        instance: Optional instance name to filter metrics
        step: Query resolution step width in duration format (e.g., 15s, 1m, 1h)
        region: The region where the cluster is located (e.g., 'us-west-2')
        o11y_id: The observability ID for the cluster (optional, will be fetched if not provided)
        tenant_id: The tenant ID (required if o11y_id is not provided)
        vendor: The cloud vendor (default: 'aws')
        
    Returns:
        Dictionary containing the query results in a Prometheus-compatible format
    """
    if not region:
        raise ValueError("Region is required for metrics query")
    
    # If o11y_id is not provided, try to fetch it using cluster_id
    if not o11y_id:
        if not tenant_id:
            raise ValueError("Either o11y_id or tenant_id must be provided")
        
        log_info(f"Looking up observability ID for cluster {cluster_id}")
        o11y_id = get_o11y_id_from_cluster_id(cluster_id, tenant_id, vendor, region)
        if not o11y_id:
            raise ValueError(f"Could not find observability ID for cluster {cluster_id}")
        log_info(f"Found observability ID: {o11y_id}")
    
    # Build the base URL
    base_url = f"https://www.ds.{region}.aws.observability.tidbcloud.com/internal/metrics/{o11y_id}/api/v1/query_range"
    
    # Initialize result structure
    result = {
        "status": "success",
        "data": {
            "resultType": "matrix",
            "result": []
        }
    }
    
    # Convert datetimes to timestamps
    start_ts = int(start_time.timestamp())
    end_ts = int(end_time.timestamp())
    
    for metric in metrics:
        try:
            # Build the query with instance filter if provided
            instance_filter = f'{{instance="{instance}"}}' if instance else ''
            query = f"{metric}{instance_filter}"
            
            # Set up query parameters
            params = {
                'query': query,
                'start': start_ts,
                'end': end_ts,
                'step': step
            }
            
            # Make the API request
            response = requests.get(base_url, params=params)
            response.raise_for_status()
            
            # Parse the response
            data = response.json()
            
            if data['status'] == 'success' and 'data' in data and 'result' in data['data']:
                # Add the results to our combined response
                for series in data['data']['result']:
                    # Ensure the metric name is included in the labels
                    if '__name__' not in series['metric']:
                        series['metric']['__name__'] = metric
                    result['data']['result'].append(series)
            
        except Exception as e:
            log_error(f"Error querying metric {metric}: {str(e)}")
            continue
    
    return result


def get_o11y_id_from_cluster_id(cluster_id: str, tenant_id: str, vendor: str, region: str) -> Optional[str]:
    """
    Get the observability ID for a given cluster ID.
    
    Args:
        cluster_id: The cluster ID to look up
        tenant_id: The tenant ID
        vendor: The cloud vendor (e.g., 'aws')
        region: The region where the cluster is located
        
    Returns:
        The observability ID if found, None otherwise
    """
    try:
        # Get auth token
        creds = get_credentials()
        auth0_token = creds.get('auth0_token')
        if not auth0_token:
            log_error("Failed to get auth token")
            return None

        # Query cluster info
        url = f"http://www.gs.{region}.aws.observability.tidbcloud.com/api/v1/clusters/"
        url += f"?bizType=tidbcloud&vendor={vendor}&region={region}&tenantID={tenant_id}&applicationID={cluster_id}"

        headers = {'authorization': f'Bearer {auth0_token}'}
        log_info(f"Querying cluster info from: {url}")
        
        resp = requests.get(url, headers=headers, timeout=30)
        resp.raise_for_status()
        
        data = resp.json()
        if not data.get('data') or not data['data'].get('clusters'):
            log_error(f"No cluster data found for cluster {cluster_id}")
            return None
            
        clusters = data['data']['clusters']
        
        # Filter out system clusters
        exclude_application_ids = ['k8s-infra', 'base', 'seed']
        exclude_project_ids = ['seed', 'base']
        
        filtered_clusters = [
            c for c in clusters
            if not (c.get('applicationID') in exclude_application_ids
                    or c.get('projectID') in exclude_project_ids)
        ]
        
        if not filtered_clusters:
            log_error(f"No valid clusters found after filtering for ID {cluster_id}")
            return None
            
        # Return the observability ID
        return filtered_clusters[0].get('id')
        
    except Exception as e:
        log_error(f"Error getting observability ID for cluster {cluster_id}: {str(e)}")
        return None


def export_metrics(
    cluster_id: str,
    start_time: datetime,
    end_time: datetime,
    output_file: str,
    instance: Optional[str] = None,
    region: Optional[str] = None,
    o11y_id: Optional[str] = None,
    tenant_id: Optional[str] = None,
    vendor: str = "aws"
) -> bool:
    """
    Export metrics in binary format using Prometheus export API.
    Exports are done in one-hour chunks for better progress control.
    
    Args:
        cluster_id: The cluster ID to export metrics for
        start_time: Start time for the export
        end_time: End time for the export
        output_file: Path to save the exported metrics
        instance: Optional instance name to filter metrics
        region: The region where the cluster is located
        o11y_id: The observability ID for the cluster
        tenant_id: The tenant ID (required if o11y_id is not provided)
        vendor: The cloud vendor (default: 'aws')
        
    Returns:
        bool: True if export was successful, False otherwise
    """
    if not region:
        raise ValueError("Region is required for metrics export")
        
    # If o11y_id is not provided, try to fetch it using cluster_id
    if not o11y_id:
        if not tenant_id:
            raise ValueError("Either o11y_id or tenant_id must be provided")
        
        log_info(f"Looking up observability ID for cluster {cluster_id}")
        o11y_id = get_o11y_id_from_cluster_id(cluster_id, tenant_id, vendor, region)
        if not o11y_id:
            raise ValueError(f"Could not find observability ID for cluster {cluster_id}")
        log_info(f"Found observability ID: {o11y_id}")
    
    # Build the base URL for export API
    base_url = f"https://www.ds.{region}.aws.observability.tidbcloud.com/internal/metrics/{o11y_id}/api/v1/export"
    
    # Prepare output directory
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Process in one-hour chunks
    chunk_size = timedelta(hours=1)
    current_start = start_time
    
    try:
        with open(output_path, 'wb') as f_out:
            while current_start < end_time:
                chunk_end = min(current_start + chunk_size, end_time)
                
                log_info(f"Exporting metrics from {current_start} to {chunk_end}")
                
                # Build the query parameters
                params = {
                    'start': current_start.isoformat() + 'Z',  # Add 'Z' for UTC
                    'end': chunk_end.isoformat() + 'Z',
                    'match[]': ['{__name__!=""}']  # Match all metrics
                }
                
                # Add instance filter if provided
                if instance:
                    params['match[]'] = [f'{{instance="{instance}"}}']
                
                # Get auth token
                creds = get_credentials()
                auth0_token = creds.get('auth0_token')
                if not auth0_token:
                    log_error("Failed to get auth token")
                    return False
                
                headers = {'authorization': f'Bearer {auth0_token}'}
                
                # Make the export request
                response = requests.get(
                    base_url,
                    params=params,
                    headers=headers,
                    stream=True,
                    timeout=300  # 5 minute timeout for each chunk
                )
                response.raise_for_status()
                
                # Write the binary data to file
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:  # filter out keep-alive new chunks
                        f_out.write(chunk)
                
                log_info(f"Exported {chunk_end - current_start} of metrics data")
                current_start = chunk_end
        
        log_info(f"Successfully exported metrics to {output_path}")
        return True
        
    except Exception as e:
        log_error(f"Error during metrics export: {str(e)}")
        # Clean up partially written file if it exists
        if output_path.exists():
            try:
                output_path.unlink()
            except Exception as cleanup_error:
                log_error(f"Failed to clean up partial export file: {str(cleanup_error)}")


def export_metrics_to_file(metrics_data: Dict[str, Any], output_format: str, output_file: str) -> None:
    """
    Export metrics data to a file in the specified format.
    
    Args:
        metrics_data: The metrics data to export
        output_format: Output format (json, csv)
        output_file: Path to the output file
    """
    try:
        if output_format == 'json':
            with open(output_file, 'w') as f:
                json.dump(metrics_data, f, indent=2)
            log_info(f"Metrics exported to {output_file} in JSON format")
            
        elif output_format == 'csv':
            # Flatten the metrics data into a CSV-friendly format
            rows = []
            for result in metrics_data.get('data', {}).get('result', []):
                metric_name = result.get('metric', {}).get('__name__', 'unknown')
                instance = result.get('metric', {}).get('instance', 'all')
                
                for value in result.get('values', []):
                    if len(value) >= 2:  # Ensure we have both timestamp and value
                        timestamp, val = value[0], value[1]
                        rows.append({
                            'timestamp': timestamp,
                            'metric': metric_name,
                            'instance': instance,
                            'value': val
                        })
            
            if rows:
                with open(output_file, 'w', newline='') as f:
                    writer = csv.DictWriter(f, fieldnames=['timestamp', 'metric', 'instance', 'value'])
                    writer.writeheader()
                    writer.writerows(rows)
                log_info(f"Metrics exported to {output_file} in CSV format")
            else:
                log_info("No metrics data to export")
        elif output_format == 'binary':
            raise ValueError("Binary export is not supported in this function. Use export_metrics() instead.")
        else:
            raise ValueError(f"Unsupported output format: {output_format}")
                
    except Exception as e:
        log_error(f"Failed to export metrics: {e}")
        raise