"""
Multiple building management module for HVAC Load Calculator.
This module provides functionality for managing multiple buildings.
"""

import streamlit as st
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
from models.building import Building, CoolingLoadResult
from controllers.enhanced_cooling_load_calculator import calculate_cooling_load
from controllers.monthly_breakdown import calculate_monthly_breakdown
from views.building_comparison import compare_buildings
from utils.data_io import export_building_to_json, import_building_from_json

class BuildingManager:
    """
    Manager for multiple buildings.
    """
    
    def __init__(self):
        """
        Initialize building manager.
        """
        # Initialize session state for buildings if not exists
        if 'buildings' not in st.session_state:
            st.session_state.buildings = {}
        
        if 'results' not in st.session_state:
            st.session_state.results = {}
        
        if 'monthly_breakdowns' not in st.session_state:
            st.session_state.monthly_breakdowns = {}
        
        if 'current_building' not in st.session_state:
            st.session_state.current_building = None
    
    def get_buildings(self) -> Dict[str, Building]:
        """
        Get all buildings.
        
        Returns:
            Dictionary of buildings {name: Building}
        """
        return st.session_state.buildings
    
    def get_results(self) -> Dict[str, CoolingLoadResult]:
        """
        Get all results.
        
        Returns:
            Dictionary of results {name: CoolingLoadResult}
        """
        return st.session_state.results
    
    def get_monthly_breakdowns(self) -> Dict[str, Dict[str, Any]]:
        """
        Get all monthly breakdowns.
        
        Returns:
            Dictionary of monthly breakdowns {name: monthly_breakdown}
        """
        return st.session_state.monthly_breakdowns
    
    def get_current_building(self) -> Optional[str]:
        """
        Get current building name.
        
        Returns:
            Current building name or None if no building is selected
        """
        return st.session_state.current_building
    
    def set_current_building(self, name: str):
        """
        Set current building.
        
        Args:
            name: Building name
        """
        st.session_state.current_building = name
    
    def add_building(self, building: Building) -> bool:
        """
        Add building to manager.
        
        Args:
            building: Building to add
            
        Returns:
            True if building was added, False if building with same name already exists
        """
        if building.settings.name in st.session_state.buildings:
            return False
        
        st.session_state.buildings[building.settings.name] = building
        return True
    
    def update_building(self, building: Building) -> bool:
        """
        Update existing building.
        
        Args:
            building: Building to update
            
        Returns:
            True if building was updated, False if building does not exist
        """
        if building.settings.name not in st.session_state.buildings:
            return False
        
        st.session_state.buildings[building.settings.name] = building
        
        # Remove old results and monthly breakdowns
        if building.settings.name in st.session_state.results:
            del st.session_state.results[building.settings.name]
        
        if building.settings.name in st.session_state.monthly_breakdowns:
            del st.session_state.monthly_breakdowns[building.settings.name]
        
        return True
    
    def remove_building(self, name: str) -> bool:
        """
        Remove building from manager.
        
        Args:
            name: Building name
            
        Returns:
            True if building was removed, False if building does not exist
        """
        if name not in st.session_state.buildings:
            return False
        
        del st.session_state.buildings[name]
        
        # Remove results and monthly breakdowns
        if name in st.session_state.results:
            del st.session_state.results[name]
        
        if name in st.session_state.monthly_breakdowns:
            del st.session_state.monthly_breakdowns[name]
        
        # Update current building if removed
        if st.session_state.current_building == name:
            if st.session_state.buildings:
                st.session_state.current_building = next(iter(st.session_state.buildings))
            else:
                st.session_state.current_building = None
        
        return True
    
    def rename_building(self, old_name: str, new_name: str) -> bool:
        """
        Rename building.
        
        Args:
            old_name: Old building name
            new_name: New building name
            
        Returns:
            True if building was renamed, False if building does not exist or new name already exists
        """
        if old_name not in st.session_state.buildings:
            return False
        
        if new_name in st.session_state.buildings:
            return False
        
        # Get building
        building = st.session_state.buildings[old_name]
        
        # Update building name
        building.settings.name = new_name
        
        # Add building with new name
        st.session_state.buildings[new_name] = building
        
        # Move results and monthly breakdowns
        if old_name in st.session_state.results:
            st.session_state.results[new_name] = st.session_state.results[old_name]
            del st.session_state.results[old_name]
        
        if old_name in st.session_state.monthly_breakdowns:
            st.session_state.monthly_breakdowns[new_name] = st.session_state.monthly_breakdowns[old_name]
            del st.session_state.monthly_breakdowns[old_name]
        
        # Remove old building
        del st.session_state.buildings[old_name]
        
        # Update current building if renamed
        if st.session_state.current_building == old_name:
            st.session_state.current_building = new_name
        
        return True
    
    def duplicate_building(self, name: str, new_name: str) -> bool:
        """
        Duplicate building.
        
        Args:
            name: Building name to duplicate
            new_name: New building name
            
        Returns:
            True if building was duplicated, False if building does not exist or new name already exists
        """
        if name not in st.session_state.buildings:
            return False
        
        if new_name in st.session_state.buildings:
            return False
        
        # Get building
        building = st.session_state.buildings[name]
        
        # Create JSON representation
        building_json = export_building_to_json(building)
        
        # Create new building from JSON
        new_building = import_building_from_json(building_json)
        
        # Update building name
        new_building.settings.name = new_name
        
        # Add new building
        st.session_state.buildings[new_name] = new_building
        
        return True
    
    def calculate_building(self, name: str) -> Tuple[CoolingLoadResult, Dict[str, Any]]:
        """
        Calculate cooling load for building.
        
        Args:
            name: Building name
            
        Returns:
            Tuple of (result, monthly_breakdown)
        """
        if name not in st.session_state.buildings:
            return None, None
        
        # Get building
        building = st.session_state.buildings[name]
        
        # Calculate cooling load
        result = calculate_cooling_load(building)
        
        # Calculate monthly breakdown
        monthly_breakdown = calculate_monthly_breakdown(building)
        
        # Store results
        st.session_state.results[name] = result
        st.session_state.monthly_breakdowns[name] = monthly_breakdown
        
        return result, monthly_breakdown
    
    def display_building_manager(self):
        """
        Display building manager UI.
        """
        st.subheader("Building Manager")
        
        # Get buildings
        buildings = self.get_buildings()
        
        if not buildings:
            st.info("No buildings available. Create a new building to get started.")
            return
        
        # Create columns
        col1, col2, col3 = st.columns([2, 1, 1])
        
        with col1:
            # Building selection
            building_names = list(buildings.keys())
            current_building = self.get_current_building()
            
            if current_building not in building_names and building_names:
                current_building = building_names[0]
            
            selected_building = st.selectbox(
                "Select Building",
                building_names,
                index=building_names.index(current_building) if current_building in building_names else 0
            )
            
            # Update current building
            self.set_current_building(selected_building)
        
        with col2:
            # Rename building
            if st.button("Rename Building"):
                new_name = st.text_input("Enter new name:", value=selected_building)
                if new_name and new_name != selected_building:
                    if self.rename_building(selected_building, new_name):
                        st.success(f"Building renamed to {new_name}")
                    else:
                        st.error(f"Failed to rename building. Name {new_name} already exists.")
        
        with col3:
            # Remove building
            if st.button("Remove Building"):
                if st.checkbox(f"Confirm removal of {selected_building}"):
                    if self.remove_building(selected_building):
                        st.success(f"Building {selected_building} removed")
                    else:
                        st.error(f"Failed to remove building {selected_building}")
        
        # Building actions
        st.subheader("Building Actions")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            # Duplicate building
            if st.button("Duplicate Building"):
                new_name = st.text_input("Enter name for duplicate:", value=f"{selected_building} (Copy)")
                if new_name:
                    if self.duplicate_building(selected_building, new_name):
                        st.success(f"Building duplicated as {new_name}")
                    else:
                        st.error(f"Failed to duplicate building. Name {new_name} already exists.")
        
        with col2:
            # Calculate building
            if st.button("Calculate Building"):
                with st.spinner(f"Calculating cooling load for {selected_building}..."):
                    result, monthly_breakdown = self.calculate_building(selected_building)
                    if result:
                        st.success(f"Calculation completed for {selected_building}")
                    else:
                        st.error(f"Failed to calculate cooling load for {selected_building}")
        
        with col3:
            # Compare buildings
            if st.button("Compare Buildings"):
                st.session_state.show_comparison = True
    
    def display_building_comparison(self):
        """
        Display building comparison UI.
        """
        # Get buildings, results, and monthly breakdowns
        buildings = self.get_buildings()
        results = self.get_results()
        monthly_breakdowns = self.get_monthly_breakdowns()
        
        # Check if we have buildings to compare
        if len(buildings) < 2:
            st.warning("At least two buildings are required for comparison. Please add more buildings.")
            return
        
        # Check if we have results for all buildings
        missing_results = [name for name in buildings if name not in results]
        if missing_results:
            st.warning(f"Missing results for buildings: {', '.join(missing_results)}. Please calculate these buildings first.")
            
            # Calculate missing results
            if st.button("Calculate Missing Results"):
                with st.spinner("Calculating missing results..."):
                    for name in missing_results:
                        self.calculate_building(name)
                    st.success("All calculations completed")
            
            return
        
        # Display comparison
        compare_buildings(buildings, results, monthly_breakdowns)
    
    def display_building_list(self):
        """
        Display building list UI.
        """
        st.subheader("Building List")
        
        # Get buildings, results, and monthly breakdowns
        buildings = self.get_buildings()
        results = self.get_results()
        monthly_breakdowns = self.get_monthly_breakdowns()
        
        if not buildings:
            st.info("No buildings available. Create a new building to get started.")
            return
        
        # Create table data
        table_data = []
        for name, building in buildings.items():
            result = results.get(name)
            monthly_breakdown = monthly_breakdowns.get(name)
            
            peak_load = result.peak_total_load if result else None
            peak_load_per_m2 = peak_load / building.settings.floor_area if peak_load else None
            
            annual_energy = monthly_breakdown.get("annual", {}).get("energy_kwh") if monthly_breakdown else None
            annual_energy_per_m2 = annual_energy / building.settings.floor_area if annual_energy else None
            
            table_data.append({
                "Building": name,
                "Location": building.location.city,
                "Floor Area (m²)": building.settings.floor_area,
                "Peak Load (kW)": peak_load / 1000 if peak_load else None,
                "Peak Load (W/m²)": peak_load_per_m2 if peak_load_per_m2 else None,
                "Annual Energy (kWh)": annual_energy if annual_energy else None,
                "Annual Energy (kWh/m²)": annual_energy_per_m2 if annual_energy_per_m2 else None,
                "Calculated": "Yes" if result else "No"
            })
        
        # Create dataframe
        df = pd.DataFrame(table_data)
        
        # Display table
        st.dataframe(df, use_container_width=True)
        
        # Add actions
        col1, col2 = st.columns(2)
        
        with col1:
            # Calculate all buildings
            if st.button("Calculate All Buildings"):
                with st.spinner("Calculating all buildings..."):
                    for name in buildings:
                        self.calculate_building(name)
                    st.success("All calculations completed")
        
        with col2:
            # Compare all buildings
            if st.button("Compare All Buildings"):
                # Check if we have results for all buildings
                missing_results = [name for name in buildings if name not in results]
                if missing_results:
                    st.warning(f"Missing results for buildings: {', '.join(missing_results)}. Please calculate these buildings first.")
                    return
                
                st.session_state.show_comparison = True