"""
PowerMCP Hugging Face Space - MCP Server Interface
Provides a Gradio interface for power system analysis tools via MCP.
"""

import json
import gradio as gr

# ============= Import from PowerMCP =============
from pandapower_tools import (
    PANDAPOWER_AVAILABLE,
    create_empty_network as pp_create_empty,
    create_test_network as pp_create_test,
    get_available_networks as pp_get_networks,
    run_power_flow as pp_run_pf,
    run_dc_power_flow as pp_run_dc,
    get_network_info as pp_get_info,
    add_bus as pp_add_bus,
    add_line as pp_add_line,
    add_load as pp_add_load,
    add_generator as pp_add_gen,
    add_ext_grid as pp_add_ext_grid,
    run_contingency_analysis as pp_contingency,
    get_available_std_types as pp_std_types,
)

from pypsa_tools import (
    PYPSA_AVAILABLE,
    create_network as pypsa_create,
    get_network_info as pypsa_get_info,
)

print(f"✓ pandapower {'loaded' if PANDAPOWER_AVAILABLE else 'not available'}")
print(f"✓ pypsa {'loaded' if PYPSA_AVAILABLE else 'not available'}")


# ============= Wrapper Functions for Gradio =============
# These wrap the PowerMCP functions to return JSON strings for Gradio display

def pandapower_create_empty_network() -> str:
    """Create an empty pandapower network."""
    return json.dumps(pp_create_empty(), indent=2, default=str)


def pandapower_create_test_network(network_type: str = "case9") -> str:
    """Create a standard IEEE test network."""
    return json.dumps(pp_create_test(network_type), indent=2, default=str)


def pandapower_run_power_flow(algorithm: str = "nr", calculate_voltage_angles: bool = True,
                               max_iteration: int = 50, tolerance_mva: float = 1e-8) -> str:
    """Run AC power flow analysis."""
    return json.dumps(pp_run_pf(algorithm, calculate_voltage_angles, 
                                max_iteration, tolerance_mva), indent=2, default=str)


def pandapower_run_dc_power_flow() -> str:
    """Run DC power flow analysis."""
    return json.dumps(pp_run_dc(), indent=2, default=str)


def pandapower_get_network_info() -> str:
    """Get network information."""
    return json.dumps(pp_get_info(), indent=2, default=str)


def pandapower_add_bus(name: str, vn_kv: float, bus_type: str = "b",
                        in_service: bool = True, max_vm_pu: float = 1.1,
                        min_vm_pu: float = 0.9) -> str:
    """Add a bus to the network."""
    return json.dumps(pp_add_bus(name, vn_kv, bus_type, in_service, 
                                 max_vm_pu, min_vm_pu), indent=2)


def pandapower_add_line(from_bus: int, to_bus: int, length_km: float,
                         std_type: str = "NAYY 4x50 SE", name: str = "") -> str:
    """Add a line to the network."""
    return json.dumps(pp_add_line(int(from_bus), int(to_bus), length_km, 
                                  std_type, name), indent=2)


def pandapower_add_load(bus: int, p_mw: float, q_mvar: float = 0.0, name: str = "") -> str:
    """Add a load to the network."""
    return json.dumps(pp_add_load(int(bus), p_mw, q_mvar, name), indent=2)


def pandapower_add_generator(bus: int, p_mw: float, vm_pu: float = 1.0,
                              name: str = "", controllable: bool = True) -> str:
    """Add a generator to the network."""
    return json.dumps(pp_add_gen(int(bus), p_mw, vm_pu, name, controllable), indent=2)


def pandapower_add_ext_grid(bus: int, vm_pu: float = 1.0, va_degree: float = 0.0,
                             name: str = "External Grid") -> str:
    """Add an external grid (slack bus) to the network."""
    return json.dumps(pp_add_ext_grid(int(bus), vm_pu, va_degree, name), indent=2)


def pandapower_run_contingency_analysis(contingency_type: str = "line",
                                         element_indices: str = "") -> str:
    """Run N-1 contingency analysis."""
    indices = None
    if element_indices.strip():
        indices = [int(x.strip()) for x in element_indices.split(",")]
    return json.dumps(pp_contingency(contingency_type, indices), indent=2, default=str)


def pandapower_get_available_std_types() -> str:
    """Get available standard types for lines and transformers."""
    return json.dumps(pp_std_types(), indent=2)


def pandapower_get_available_networks() -> str:
    """Get a list of all available pandapower test networks."""
    return json.dumps(pp_get_networks(), indent=2)


def pypsa_create_network(name: str = "PyPSA Network") -> str:
    """Create a new PyPSA network."""
    return json.dumps(pypsa_create(name), indent=2)


def pypsa_get_network_info() -> str:
    """Get PyPSA network information."""
    return json.dumps(pypsa_get_info(), indent=2)


def get_available_tools() -> str:
    """Get a list of all available tools."""
    tools = {
        "pandapower_tools": {
            "available": PANDAPOWER_AVAILABLE,
            "tools": [
                "pandapower_create_empty_network",
                "pandapower_create_test_network",
                "pandapower_get_available_networks",
                "pandapower_run_power_flow",
                "pandapower_run_dc_power_flow",
                "pandapower_get_network_info",
                "pandapower_add_bus",
                "pandapower_add_line",
                "pandapower_add_load",
                "pandapower_add_generator",
                "pandapower_add_ext_grid",
                "pandapower_run_contingency_analysis",
                "pandapower_get_available_std_types"
            ]
        },
        "pypsa_tools": {
            "available": PYPSA_AVAILABLE,
            "tools": [
                "pypsa_create_network",
                "pypsa_get_network_info"
            ]
        }
    }
    return json.dumps(tools, indent=2)


# ============= Gradio Interfaces =============

iface_pp_create_empty = gr.Interface(
    fn=pandapower_create_empty_network,
    inputs=[],
    outputs=gr.Textbox(label="Result"),
    title="Create Empty Pandapower Network",
    description="Create an empty pandapower network"
)

iface_pp_create_test = gr.Interface(
    fn=pandapower_create_test_network,
    inputs=[gr.Textbox(label="Network Type", value="case9",
                       info="Use pandapower_get_available_networks to see all options. Examples: case9, case14, case30, case39, case57, case118, create_cigre_network_hv, iceland, GBnetwork, etc.")],
    outputs=gr.Textbox(label="Result"),
    title="Create Test Network",
    description="Create a standard IEEE test network or other built-in pandapower network"
)

iface_pp_get_networks = gr.Interface(
    fn=pandapower_get_available_networks,
    inputs=[],
    outputs=gr.Textbox(label="Result"),
    title="Get Available Networks",
    description="Get a list of all available pandapower test networks"
)

iface_pp_run_pf = gr.Interface(
    fn=pandapower_run_power_flow,
    inputs=[
        gr.Textbox(label="Algorithm", value="nr", info="nr, bfsw, gs, fdbx, fdxb"),
        gr.Checkbox(label="Calculate Voltage Angles", value=True),
        gr.Number(label="Max Iterations", value=50),
        gr.Number(label="Tolerance (MVA)", value=1e-8)
    ],
    outputs=gr.Textbox(label="Result"),
    title="Run AC Power Flow",
    description="Run AC power flow analysis"
)

iface_pp_run_dc = gr.Interface(
    fn=pandapower_run_dc_power_flow,
    inputs=[],
    outputs=gr.Textbox(label="Result"),
    title="Run DC Power Flow",
    description="Run DC power flow analysis"
)

iface_pp_get_info = gr.Interface(
    fn=pandapower_get_network_info,
    inputs=[],
    outputs=gr.Textbox(label="Result"),
    title="Get Network Info",
    description="Get network information"
)

iface_pp_add_bus = gr.Interface(
    fn=pandapower_add_bus,
    inputs=[
        gr.Textbox(label="Bus Name"),
        gr.Number(label="Nominal Voltage (kV)", value=20.0),
        gr.Textbox(label="Bus Type", value="b"),
        gr.Checkbox(label="In Service", value=True),
        gr.Number(label="Max Voltage (pu)", value=1.1),
        gr.Number(label="Min Voltage (pu)", value=0.9)
    ],
    outputs=gr.Textbox(label="Result"),
    title="Add Bus",
    description="Add a bus to the network"
)

iface_pp_add_line = gr.Interface(
    fn=pandapower_add_line,
    inputs=[
        gr.Number(label="From Bus Index", precision=0),
        gr.Number(label="To Bus Index", precision=0),
        gr.Number(label="Length (km)", value=1.0),
        gr.Textbox(label="Standard Type", value="NAYY 4x50 SE"),
        gr.Textbox(label="Name", value="")
    ],
    outputs=gr.Textbox(label="Result"),
    title="Add Line",
    description="Add a line to the network"
)

iface_pp_add_load = gr.Interface(
    fn=pandapower_add_load,
    inputs=[
        gr.Number(label="Bus Index", precision=0),
        gr.Number(label="Active Power (MW)", value=1.0),
        gr.Number(label="Reactive Power (Mvar)", value=0.0),
        gr.Textbox(label="Name", value="")
    ],
    outputs=gr.Textbox(label="Result"),
    title="Add Load",
    description="Add a load to the network"
)

iface_pp_add_gen = gr.Interface(
    fn=pandapower_add_generator,
    inputs=[
        gr.Number(label="Bus Index", precision=0),
        gr.Number(label="Active Power (MW)", value=10.0),
        gr.Number(label="Voltage Setpoint (pu)", value=1.0),
        gr.Textbox(label="Name", value=""),
        gr.Checkbox(label="Controllable", value=True)
    ],
    outputs=gr.Textbox(label="Result"),
    title="Add Generator",
    description="Add a generator to the network"
)

iface_pp_add_ext_grid = gr.Interface(
    fn=pandapower_add_ext_grid,
    inputs=[
        gr.Number(label="Bus Index", precision=0),
        gr.Number(label="Voltage Magnitude (pu)", value=1.0),
        gr.Number(label="Voltage Angle (degrees)", value=0.0),
        gr.Textbox(label="Name", value="External Grid")
    ],
    outputs=gr.Textbox(label="Result"),
    title="Add External Grid",
    description="Add an external grid (slack bus)"
)

iface_pp_contingency = gr.Interface(
    fn=pandapower_run_contingency_analysis,
    inputs=[
        gr.Textbox(label="Contingency Type", value="line", info="line, trafo, or gen"),
        gr.Textbox(label="Element Indices", value="", info="Comma-separated, empty for all")
    ],
    outputs=gr.Textbox(label="Result"),
    title="Run Contingency Analysis",
    description="Run N-1 contingency analysis"
)

iface_pp_std_types = gr.Interface(
    fn=pandapower_get_available_std_types,
    inputs=[],
    outputs=gr.Textbox(label="Result"),
    title="Get Standard Types",
    description="Get available standard types"
)

iface_pypsa_create = gr.Interface(
    fn=pypsa_create_network,
    inputs=[gr.Textbox(label="Network Name", value="PyPSA Network")],
    outputs=gr.Textbox(label="Result"),
    title="Create PyPSA Network",
    description="Create a new PyPSA network"
)

iface_pypsa_info = gr.Interface(
    fn=pypsa_get_network_info,
    inputs=[],
    outputs=gr.Textbox(label="Result"),
    title="Get PyPSA Network Info",
    description="Get PyPSA network information"
)

iface_list_tools = gr.Interface(
    fn=get_available_tools,
    inputs=[],
    outputs=gr.Textbox(label="Available Tools"),
    title="List Available Tools",
    description="List all available tools"
)


# ============= Main Gradio App =============

with gr.Blocks(title="PowerMCP Server") as demo:
    gr.Markdown("""
    # ⚡ PowerMCP - Power System Analysis MCP Server
    
    This Space exposes power system analysis tools via the **Model Context Protocol (MCP)**.
    
    **Use this Space:**
    - 🤖 **As MCP tools** in Claude Desktop, Cursor, or other MCP clients
    - 🖥️ **Web interface** below for manual testing
    
    **Available Tools:** pandapower (power flow, contingency analysis), PyPSA (optimization)
    """)
    
    with gr.Tab("🔌 Quick Start"):
        gr.Markdown("### Create and Analyze a Test Network")
        with gr.Row():
            with gr.Column():
                test_network_type = gr.Dropdown(
                    choices=["case9", "case14", "case30", "case39", "case57", "case118"],
                    value="case9",
                    label="Select Test Network"
                )
                create_btn = gr.Button("🔧 Create Network", variant="primary")
                create_output = gr.Code(label="Creation Result", language="json", lines=8)
            
            with gr.Column():
                run_pf_btn = gr.Button("⚡ Run Power Flow", variant="primary")
                pf_output = gr.Code(label="Power Flow Results", language="json", lines=8)
            
            with gr.Column():
                contingency_btn = gr.Button("🔍 Run Contingency Analysis", variant="secondary")
                contingency_output = gr.Code(label="Contingency Results", language="json", lines=8)
        
        info_btn = gr.Button("ℹ️ Get Network Info", variant="secondary")
        info_output = gr.Code(label="Network Info", language="json", lines=10)
        
        create_btn.click(fn=pandapower_create_test_network, inputs=[test_network_type], outputs=[create_output])
        run_pf_btn.click(fn=pandapower_run_power_flow, inputs=[], outputs=[pf_output])
        contingency_btn.click(fn=pandapower_run_contingency_analysis, inputs=[], outputs=[contingency_output])
        info_btn.click(fn=pandapower_get_network_info, inputs=[], outputs=[info_output])
    
    with gr.Tab("🔧 Pandapower Tools"):
        gr.TabbedInterface(
            [iface_pp_create_empty, iface_pp_create_test, iface_pp_get_networks, iface_pp_run_pf, iface_pp_run_dc,
             iface_pp_get_info, iface_pp_add_bus, iface_pp_add_line, iface_pp_add_load,
             iface_pp_add_gen, iface_pp_add_ext_grid, iface_pp_contingency, iface_pp_std_types],
            ["Create Empty", "Create Test", "List Networks", "AC Power Flow", "DC Power Flow",
             "Get Info", "Add Bus", "Add Line", "Add Load", "Add Gen", "Add Ext Grid",
             "Contingency", "Std Types"]
        )
    
    with gr.Tab("⚡ PyPSA Tools"):
        gr.TabbedInterface(
            [iface_pypsa_create, iface_pypsa_info],
            ["Create Network", "Get Info"]
        )
    
    with gr.Tab("📚 Documentation"):
        gr.Markdown("""
        ## MCP Server Usage
        
        ### Connect via MCP Client
        
        Add this Space to your MCP client configuration:
        
        ```json
        {
          "mcpServers": {
            "powermcp": {
              "url": "https://your-space-url.hf.space/gradio_api/mcp/sse"
            }
          }
        }
        ```
        
        ### Available Pandapower Tools
        
        | Tool | Description |
        |------|-------------|
        | `pandapower_create_empty_network` | Create an empty network |
        | `pandapower_create_test_network` | Create IEEE test networks (all pandapower networks supported) |
        | `pandapower_get_available_networks` | List all available test networks |
        | `pandapower_run_power_flow` | Run AC power flow |
        | `pandapower_run_dc_power_flow` | Run DC power flow |
        | `pandapower_get_network_info` | Get network information |
        | `pandapower_add_bus` | Add a bus |
        | `pandapower_add_line` | Add a line |
        | `pandapower_add_load` | Add a load |
        | `pandapower_add_generator` | Add a generator |
        | `pandapower_add_ext_grid` | Add external grid |
        | `pandapower_run_contingency_analysis` | Run N-1 contingency |
        | `pandapower_get_available_std_types` | Get standard types |
        
        ### Available PyPSA Tools
        
        | Tool | Description |
        |------|-------------|
        | `pypsa_create_network` | Create a PyPSA network |
        | `pypsa_get_network_info` | Get network info |
        """)
    
    with gr.Tab("🛠️ Available Tools"):
        gr.Interface(
            fn=get_available_tools,
            inputs=[],
            outputs=gr.Code(label="Tools", language="json"),
            title="Available Tools",
            description="View all available power system tools"
        )


def main():
    """Entry point for powermcp-app command"""
    # Try MCP server mode if available (requires gradio[mcp])
    try:
        demo.launch(mcp_server=True, share=False)
    except TypeError:
        # Fallback for older Gradio versions without MCP support
        demo.launch(share=False)


if __name__ == "__main__":
    main()