import gradio as gr
import xml.etree.ElementTree as ET
import io

def get_total_capacity_from_xml(xml_file_path, substation_names_str, profile, flexibility_level):
    """
    Calculates the total capacity for specified substations, profile, and flexibility level
    from an uploaded XML file's content.

    Args:
        xml_file_path (str): The file path (string) of the temporary uploaded XML file.
        substation_names_str (str): A comma-separated string of substation names.
        profile (str): The profile type (e.g., "pv", "load", "battery").
        flexibility_level (str): The flexibility level (e.g., "0%_flex", "5%_flex").

    Returns:
        str: A message indicating the total combined capacity or an error.
    """
    if xml_file_path is None:
        return "Error: No XML file uploaded."

    try:
        # Read the content from the temporary file path provided by Gradio
        with open(xml_file_path, 'r', encoding='utf-8') as f:
            xml_file_content = f.read()
    except Exception as e:
        return f"Error reading XML file: {e}"

    if not xml_file_content:
        return "Error: Uploaded XML file is empty or could not be read."

    substation_names = [s.strip() for s in substation_names_str.split(',') if s.strip()]
    if not substation_names:
        return "Error: Please enter at least one substation name."

    total_capacity = 0.0
    
    # Define namespaces used in the XML
    namespaces = {
        'default': 'http://schemas.datacontract.org/2004/07/Elia.EliaBeControls.WebApi.Interface.Dto.Hcm',
        'd5p1': 'http://schemas.microsoft.com/2003/10/Serialization/Arrays'
    }

    try:
        # Parse the XML content from the string
        root = ET.parse(io.StringIO(xml_file_content)).getroot()
    except ET.ParseError as e:
        return f"Error parsing XML: The file might be malformed or not valid XML. Details: {e}"
    except Exception as e:
        return f"An unexpected error occurred during XML parsing: {e}"

    found_any_matching_data = False

    # Iterate through each SiteDetailDto element
    for site_detail_dto in root.findall('default:SiteDetailDto', namespaces):
        site_name_element = site_detail_dto.find('default:Site', namespaces)

        # Check if the site name element exists and its text is in our list of substations
        if site_name_element is not None and site_name_element.text in substation_names:
            # Iterate through each SubstationDetailDto within the site
            for substation_detail_dto in site_detail_dto.findall('.//default:SubstationDetailDto', namespaces):
                # Iterate through profiles
                for profile_kv in substation_detail_dto.findall('default:Profiles/d5p1:KeyValueOfstringArrayOfProfileInfoDtoOJezfHCJ', namespaces):
                    key_element = profile_kv.find('d5p1:Key', namespaces)

                    # Check if the profile key matches the requested profile
                    if key_element is not None and key_element.text == profile:
                        # Iterate through flexibility levels within the profile
                        for profile_info_dto in profile_kv.findall('d5p1:Value/default:ProfileInfoDto', namespaces):
                            flexibility_level_element = profile_info_dto.find('default:FlexibilityLevel', namespaces)

                            # Check if the flexibility level matches the requested level
                            if flexibility_level_element is not None and flexibility_level_element.text == flexibility_level:
                                capacity_element = profile_info_dto.find('default:FlexibilityDetail/default:Capacity', namespaces)

                                # Check if the capacity element exists and has text
                                if capacity_element is not None and capacity_element.text:
                                    try:
                                        total_capacity += float(capacity_element.text)
                                        found_any_matching_data = True
                                    except ValueError:
                                        # Log or handle cases where capacity might not be a valid float
                                        print(f"Warning: Could not convert capacity '{capacity_element.text}' to float for {site_name_element.text}, Profile: {profile}, Flex: {flexibility_level}")
                                        pass # Continue processing other elements

    if not found_any_matching_data and total_capacity == 0.0:
        return f"No matching data found for substations: {substation_names_str}, profile: '{profile}', and flexibility: '{flexibility_level}'. Check your inputs and the XML structure."
    
    return f"The total capacity for the selected criteria is: {total_capacity}"

# Gradio Interface setup
# Input components
# Changed type="filepath" to ensure a string path is passed to the function
xml_upload = gr.File(label="Upload XML File", file_types=[".xml"], type="filepath")
substation_input = gr.Textbox(label="Enter Substation(s) (comma-separated, e.g., AALST,AALTR)", placeholder="AALST, AALTR")
profile_dropdown = gr.Dropdown(choices=["pv", "load", "battery"], label="Select Profile", value="pv")
flexibility_dropdown = gr.Dropdown(choices=["0%_flex", "5%_flex", "10%_flex", "20%_flex"], label="Select Flexibility Level", value="5%_flex")

# Define the Gradio interface
iface = gr.Interface(
    fn=get_total_capacity_from_xml,
    inputs=[xml_upload, substation_input, profile_dropdown, flexibility_dropdown],
    outputs="text",
    title="Substation Capacity Calculator from XML",
    description="Upload an XML file, then select substations, profile, and flexibility level to get the total capacity. This tool performs basic structural checks on the XML."
)

# Launch the Gradio app
iface.launch(share=True)