#!/usr/bin/env python3
"""
National Grid Real-Time Data Monitor
Fetches load, fuel mixture, and other real-time data for a configurable GSP
Pushes metrics to VictoriaMetrics Prometheus endpoint
"""

import requests
import json
import sys
from datetime import datetime
from typing import Dict, Any, Optional, List

# GSP ID mapping (based on Carbon Intensity API regional IDs)
GSP_MAPPING = {
    "Bishops Wood": 11,  # West Midlands
    "North Scotland": 1,
    "South Scotland": 2,
    "North West England": 3,
    "North East England": 4,
    "Yorkshire": 5,
    "North Wales": 6,
    "South Wales": 7,
    "West Midlands": 11,
    "East Midlands": 12,
    "East England": 13,
    "South West England": 14,
    "South England": 15,
    "London": 16,
    "South East England": 17,
}

class NationalGridMonitor:
    """Monitor for National Grid real-time data"""
    
    BASE_URL = "https://api.carbonintensity.org.uk"
    
    def __init__(self, gsp_name: str = "Bishops Wood", victoriametrics_url: str = None):
        self.gsp_name = gsp_name
        self.gsp_id = GSP_MAPPING.get(gsp_name)
        self.victoriametrics_url = victoriametrics_url
        
        if self.gsp_id is None:
            print(f"Warning: GSP '{gsp_name}' not found in mapping. Available GSPs:")
            for name in sorted(GSP_MAPPING.keys()):
                print(f"  - {name}")
            sys.exit(1)
    
    def get_regional_data(self) -> Optional[Dict[str, Any]]:
        """Fetch regional data including generation mix and intensity"""
        try:
            url = f"{self.BASE_URL}/regional/regionid/{self.gsp_id}"
            params = {"resolution": "minute"}
            response = requests.get(url, params=params, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error fetching regional data: {e}")
            return None
    
    def get_national_stats(self) -> Optional[Dict[str, Any]]:
        """Fetch national statistics including demand"""
        try:
            url = f"{self.BASE_URL}/intensity"
            params = {"resolution": "minute"}
            response = requests.get(url, params=params, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error fetching national stats: {e}")
            return None
    
    def get_generation_mix(self) -> Optional[Dict[str, Any]]:
        """Fetch current generation mix"""
        try:
            url = f"{self.BASE_URL}/generation"
            params = {"resolution": "minute"}
            response = requests.get(url, params=params, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error fetching generation mix: {e}")
            return None
    
    def get_demand_data(self) -> Optional[Dict[str, Any]]:
        """Fetch actual demand data from Elexon BMRS API (no API key required)"""
        try:
            # Elexon BMRS API provides actual demand in MW
            url = "https://data.elexon.co.uk/bmrs/api/v1/demand/outturn/summary"
            params = {"resolution": "minute"}
            response = requests.get(url, params=params, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error fetching demand data: {e}")
            return None
    
    def collect_data(self) -> Dict[str, Any]:
        """Collect all available data and format for output"""
        
        print(f"Fetching data for GSP: {self.gsp_name} (ID: {self.gsp_id})...")
        
        regional_data = self.get_regional_data()
        national_stats = self.get_national_stats()
        generation_mix = self.get_generation_mix()
        demand_data = self.get_demand_data()
        
        # Build comprehensive data structure
        output = {
            "gsp_name": self.gsp_name,
            "gsp_id": self.gsp_id,
            "timestamp": datetime.now(datetime.UTC if hasattr(datetime, 'UTC') else None).replace(tzinfo=None).isoformat() + "Z",
            "collection_time": datetime.now().isoformat(),
        }
        
        # Regional data
        if regional_data and "data" in regional_data:
            data_entry = regional_data["data"][0] if regional_data["data"] else {}
            output["regional"] = {
                "from": data_entry.get("from"),
                "to": data_entry.get("to"),
                "intensity": data_entry.get("intensity", {}),
                "generation_mix": data_entry.get("generationmix", [])
            }
        
        # National statistics
        if national_stats and "data" in national_stats:
            data_entry = national_stats["data"][0] if national_stats["data"] else {}
            output["national"] = {
                "from": data_entry.get("from"),
                "to": data_entry.get("to"),
                "intensity": data_entry.get("intensity", {})
            }
        
        # National generation mix
        if generation_mix and "data" in generation_mix:
            data_entry = generation_mix["data"]
            output["generation"] = {
                "from": data_entry.get("from"),
                "to": data_entry.get("to"),
                "mix": data_entry.get("generationmix", [])
            }
        
        # Actual demand from Elexon BMRS API (in MW)
        if demand_data and isinstance(demand_data, list) and len(demand_data) > 0:
            # Get the most recent demand reading
            latest_demand = demand_data[-1]
            output["demand"] = {
                "timestamp": latest_demand.get("startTime"),
                "demand_mw": latest_demand.get("demand"),
                "record_type": latest_demand.get("recordType"),
                "source": "Elexon BMRS API"
            }
        else:
            output["demand"] = {
                "demand_mw": None,
                "note": "Demand data unavailable"
            }
        
        return output
    
    def save_to_file(self, data: Dict[str, Any]) -> None:
        """Save data to JSON file"""
        filename = f"{self.gsp_name.replace(' ', '_')}.json"
        
        with open(filename, 'w') as f:
            json.dump(data, f, indent=2)
        
        print(f"Data saved to: {filename}")
    
    def format_prometheus_metrics(self, data: Dict[str, Any]) -> str:
        """Format data as Prometheus exposition format for VictoriaMetrics"""
        metrics = []
        gsp_label = self.gsp_name.replace(' ', '_').replace('-', '_')
        timestamp_ms = int(datetime.now().timestamp() * 1000)
        
        # Carbon intensity
        if "national" in data and "intensity" in data["national"]:
            intensity = data["national"]["intensity"]
            if "actual" in intensity and intensity["actual"] is not None:
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="carbon_intensity"}} {intensity["actual"]} {timestamp_ms}'
                )
            if "forecast" in intensity and intensity["forecast"] is not None:
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="carbon_intensity_forecast"}} {intensity["forecast"]} {timestamp_ms}'
                )
        
        # Generation mix percentages
        if "generation" in data and "mix" in data["generation"]:
            for fuel_data in data["generation"]["mix"]:
                fuel = fuel_data["fuel"]
                perc = fuel_data["perc"]
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="{fuel}"}} {perc} {timestamp_ms}'
                )
        
        # Regional generation mix (if available)
        if "regional" in data and "generation_mix" in data["regional"]:
            for fuel_data in data["regional"]["generation_mix"]:
                fuel = fuel_data["fuel"]
                perc = fuel_data["perc"]
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="regional_{fuel}"}} {perc} {timestamp_ms}'
                )
        
        # Regional intensity (if available)
        if "regional" in data and "intensity" in data["regional"]:
            intensity = data["regional"]["intensity"]
            if "forecast" in intensity and intensity["forecast"] is not None:
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="regional_carbon_intensity"}} {intensity["forecast"]} {timestamp_ms}'
                )
        
        # Actual demand in MW from Elexon
        if "demand" in data and "demand_mw" in data["demand"]:
            demand_mw = data["demand"]["demand_mw"]
            if demand_mw is not None:
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="demand_mw"}} {demand_mw} {timestamp_ms}'
                )
        
        # Carbon intensity as a load indicator
        if "national" in data and "intensity" in data["national"]:
            actual_intensity = data["national"]["intensity"].get("actual")
            if actual_intensity is not None:
                metrics.append(
                    f'national_grid{{gsp="{gsp_label}",metric="load_indicator"}} {actual_intensity} {timestamp_ms}'
                )
        
        return "\n".join(metrics) + "\n"
    
    def push_to_victoriametrics(self, data: Dict[str, Any]) -> bool:
        """Push metrics to VictoriaMetrics using Prometheus exposition format"""
        if not self.victoriametrics_url:
            return False
        
        try:
            metrics_text = self.format_prometheus_metrics(data)
            
            # VictoriaMetrics supports multiple ingestion endpoints:
            # 1. /api/v1/import/prometheus - VictoriaMetrics native (Prometheus text format)
            # 2. /api/v1/write - Prometheus remote_write protocol (requires protobuf)
            # Using native import endpoint for simplicity
            url = f"{self.victoriametrics_url}/api/v1/import/prometheus"
            
            headers = {
                "Content-Type": "text/plain"
            }
            
            response = requests.post(url, data=metrics_text, headers=headers, timeout=10)
            response.raise_for_status()
            
            print(f"✓ Pushed {len(metrics_text.strip().split(chr(10)))} metrics to VictoriaMetrics: {url}")
            return True
            
        except requests.exceptions.RequestException as e:
            print(f"✗ Error pushing to VictoriaMetrics: {e}")
            return False
    
    def run(self) -> None:
        """Main execution method"""
        data = self.collect_data()
        self.save_to_file(data)
        
        # Push to VictoriaMetrics if configured
        if self.victoriametrics_url:
            self.push_to_victoriametrics(data)
        
        # Print summary
        print("\n=== Data Summary ===")
        
        # National intensity
        if "national" in data and "intensity" in data["national"]:
            intensity = data["national"]["intensity"]
            print(f"National Carbon Intensity: {intensity.get('actual', 'N/A')} gCO2/kWh (forecast: {intensity.get('forecast', 'N/A')})")
            print(f"Index: {intensity.get('index', 'N/A')}")
        
        # Regional intensity
        if "regional" in data and "intensity" in data["regional"] and data["regional"]["intensity"]:
            print(f"Regional Intensity: {data['regional']['intensity'].get('forecast')} gCO2/kWh")
            print(f"Regional Index: {data['regional']['intensity'].get('index')}")
        
        # National generation mix
        if "generation" in data:
            print("\nNational Generation Mix:")
            for gen in data["generation"]["mix"]:
                print(f"  {gen['fuel']:12s}: {gen['perc']:5.1f}%")
        
        # Regional generation mix (if different/available)
        if "regional" in data and data["regional"].get("generation_mix"):
            print("\nRegional Generation Mix:")
            for gen in data["regional"]["generation_mix"]:
                print(f"  {gen['fuel']:12s}: {gen['perc']:5.1f}%")
        
        # Demand info
        if "demand" in data:
            demand_mw = data['demand'].get('demand_mw', 'N/A')
            if demand_mw != 'N/A' and demand_mw is not None:
                print(f"\nNational Demand: {demand_mw:,.0f} MW")
                print(f"Source: {data['demand'].get('source', 'Unknown')}")
                if "timestamp" in data['demand']:
                    print(f"Timestamp: {data['demand']['timestamp']}")
            else:
                print(f"\nDemand: {data['demand'].get('note', 'N/A')}")


def main():
    """Main entry point"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Fetch National Grid real-time data for a GSP and push to VictoriaMetrics"
    )
    parser.add_argument(
        "--gsp",
        default="Bishops Wood",
        help="Grid Supply Point name (default: Bishops Wood)"
    )
    parser.add_argument(
        "--victoriametrics",
        default="http://localhost:8428",
        help="VictoriaMetrics URL (default: http://localhost:8428)"
    )
    parser.add_argument(
        "--no-push",
        action="store_true",
        help="Don't push to VictoriaMetrics, only save to JSON"
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List available GSP names"
    )
    
    args = parser.parse_args()
    
    if args.list:
        print("Available GSP names:")
        for name in sorted(GSP_MAPPING.keys()):
            print(f"  - {name}")
        sys.exit(0)
    
    vm_url = None if args.no_push else args.victoriametrics
    monitor = NationalGridMonitor(gsp_name=args.gsp, victoriametrics_url=vm_url)
    monitor.run()


if __name__ == "__main__":
    main()