from typing import defaultdict, List, Optional, Union
from uuid import UUID
import logging
import httpx
from fastapi import HTTPException, status
from sqlalchemy import Delete, Select, func, desc, and_, text
from sqlalchemy.orm import selectinload

from src.auth.service import CurrentUser
from src.config import DEFAULT_PROJECT_NAME, RELIABILITY_SERVICE_API
from src.aeros_utils import aeros_post
from src.database.core import CollectorDbSession, DbSession
from src.database.service import search_filter_sort_paginate
from .model import AerosEquipment, AerosEquipmentDetail, MasterEquipment, AerosEquipmentGroup, ReliabilityPredictNonRepairable
from .schema import EquipmentConfiguration
from src.aeros_project.model import AerosProject
from src.aeros_simulation.model import AerosNode
import asyncio
import re
import requests
import json
import pandas as pd
from src.aeros_simulation.service import get_aeros_schematic_by_name

# Aeros session is managed in src.aeros_utils
log = logging.getLogger()

async def get_project(*, db_session: DbSession):
    stmt = Select(AerosProject).order_by(desc(AerosProject.updated_at)).limit(1)
    result = await db_session.execute(stmt)
    found_record = result.scalar_one_or_none()

    return found_record

async def get_all(*, common):
    """Returns all documents."""
    query = Select(AerosEquipment).options(
        selectinload(AerosEquipment.master_equipment)
    )
    results = await search_filter_sort_paginate(model=query, **common)
    reg_nodes = [node.node_name for node in results["items"]]
    equipment_data = {
        node.node_name: node for node in results["items"]
    }

    project = await get_project(db_session=common["db_session"])

    updateNodeReq = {"projectName": project.project_name , "equipmentNames": reg_nodes}

    try:
        response = await aeros_post(
            "/api/UpdateDisParams/GetUpdatedNodeDistributions",
            json=updateNodeReq,
            headers={"Content-Type": "application/json"},
        )
        
        response.raise_for_status()
        res = response.json()
        # if not res.get("status"):
        #     raise HTTPException(
        #         status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=res.get("message")
        #     )

    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    results["items"] = [
        {"AerosData": data, "MasterData": equipment_data.get(data["equipmentName"]) } for data in res
    ]

    return results

async def get_equipment_by_location_tag(*, db_session: DbSession, location_tag: str):
    query = (
        Select(AerosEquipment)
        .where(AerosEquipment.location_tag == location_tag)
        .options(selectinload(AerosEquipment.aeros_equipment_details))
    )



async def get_by_id(*, db_session: DbSession, id: UUID):
    query = (
        Select(AerosEquipment)
        .where(AerosEquipment.id == id)
        .options(selectinload(AerosEquipment.aeros_equipment_details))
    )
    result = await db_session.execute(query)
    aerosEquipmentResult = result.scalar()

    if not aerosEquipmentResult:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="AerosEquipment not found"
        )

    aerosNodeReq = {
        "projectName": "ParallelNode",
        "equipmentName": [aerosEquipmentResult.node_name],
    }

    try:
        response = await aeros_post(
            "/api/UpdateDisParams/GetUpdatedNodeDistributions",
            json=aerosNodeReq,
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
        aerosEquipmentData = response.json()

    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )

    return aerosEquipmentResult, aerosEquipmentData


async def get_aeros_equipment_by_location_tag(*, location_tag: Union[str, List[str]], project_name):
    if isinstance(location_tag, str):
        location_tag = [location_tag]
    else:
        location_tag = location_tag

    aerosNodeReq = {
        "projectName": project_name,
        "equipmentNames": location_tag,
    }

    try:
        response = await aeros_post(
            "/api/UpdateDisParams/GetUpdatedNodeDistributions",
            json=aerosNodeReq,
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()
        aerosEquipmentData = response.json()
    

        if not aerosEquipmentData:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND, detail="AerosEquipment not found"
            )
        
        response.raise_for_status()

        df = pd.DataFrame(aerosEquipmentData)

        return df.drop_duplicates().to_dict('records')

    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )


async def update_node(
    *, db_session: DbSession, equipment_nodes: List[dict], project_name: str
):
    updateNodeReq = {"projectName": project_name, "regNodeInputs": equipment_nodes}


    try:
        response = await aeros_post(
            "/api/UpdateDisParams/UpdateEquipmentDistributions",
            json=updateNodeReq,
            headers={"Content-Type": "application/json"},
        )

        result = response.json()

        return result

    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )


async def save_default_equipment(*, db_session: DbSession, project_name: str):
    equipmments = Select(MasterEquipment).where(
        MasterEquipment.location_tag.isnot(None)
    )
    equipment_nodes = await db_session.execute(equipmments)

    reg_nodes = [node.location_tag for node in equipment_nodes.scalars().all()]

    updateNodeReq = {"projectName": project_name, "equipmentNames": reg_nodes}

    # Delete old data
    query = Delete(AerosEquipment)
    await db_session.execute(query)

    try:
        response = await aeros_post(
            "/api/UpdateDisParams/GetUpdatedNodeDistributions",
            json=updateNodeReq,
            headers={"Content-Type": "application/json"},
        )
        response.raise_for_status()

        results = response.json()

        nodes = []

        # raise Exception(results)

        # save to db
        for equipment in results:
            node = AerosEquipment(
                node_name=equipment["equipmentName"],
                location_tag=equipment["equipmentName"],
            )

            nodes.append(node)

        db_session.add_all(nodes)
        await db_session.commit()

        return results
    except Exception as e:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )


async def get_asset_batch(location_tags: List[str], nr_location_tags: List[str],
                    db_session: DbSession,
                    base_url: str = RELIABILITY_SERVICE_API,
                   timeout: int = 600)-> dict:
    """
    Get asset batch data from database tables.

    Args:
        location_tags: List of location tag strings for repairable equipment
        nr_location_tags: List of location tag strings for non-repairable equipment
        db_session: Database session
        base_url: Base URL for the API (kept for compatibility)
        timeout: Request timeout in seconds (kept for compatibility)

    Returns:
        Dictionary with response data
    """
    from sqlalchemy import text
    import json
    
    results = defaultdict(dict)

    # Query repairable equipment data
    if location_tags:
        repairable_query = text("""
            SELECT location_tag, parameters
            FROM rp_repairable_results
            WHERE location_tag = ANY(:location_tags)
        """)
        
        repairable_result = await db_session.execute(repairable_query, {"location_tags": location_tags})
        
        for row in repairable_result:
            location_tag = row[0]
            raw_parameters = row[1]
            parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {}
            
            try:
                distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=True)
                
                results[location_tag]["cmDisType"] = "Normal"
                results[location_tag]["cmDisP1"] = 6
                results[location_tag]["cmDisP2"] = 3
                results[location_tag]["relDisType"] = distribution
                results[location_tag]["relDisP1"] = reldisp1
                results[location_tag]["relDisP2"] = reldisp2
                results[location_tag]["parameters"] = parameters
            except Exception as e:
                raise Exception(f"Error processing repairable item {location_tag}: {e}")

    # Query non-repairable equipment data
    if nr_location_tags:
        non_repairable_query = text("""
            SELECT location_tag, parameters, best_distribution
            FROM rp_non_repairable_results
            WHERE location_tag = ANY(:nr_location_tags)
        """)
        
        non_repairable_result = await db_session.execute(non_repairable_query, {"nr_location_tags": nr_location_tags})
        
        for row in non_repairable_result:
            location_tag = row[0]
            raw_parameters = row[1]
            parameters = raw_parameters if isinstance(raw_parameters, dict) else json.loads(raw_parameters) if raw_parameters else {}
            best_distribution = row[2]
            
            try:
                distribution, reldisp1, reldisp2 = get_distribution_from_db(parameters, is_repairable=False, best_distribution=best_distribution)
                
                results[location_tag]["cmDisType"] = "Normal"
                results[location_tag]["cmDisP1"] = 6
                results[location_tag]["cmDisP2"] = 3
                results[location_tag]["relDisType"] = distribution
                results[location_tag]["relDisP1"] = reldisp1
                results[location_tag]["relDisP2"] = reldisp2
                results[location_tag]["parameters"] = parameters
            except Exception as e:
                raise Exception(f"Error processing non-repairable item {location_tag}: {e}")

    return results


def get_distribution_from_db(parameters: dict, is_repairable: bool = True, best_distribution: str = None):
    """
    Map distribution parameters from database to the expected format.
    
    Args:
        parameters: Parameters JSON from database
        is_repairable: Whether this is repairable equipment
        best_distribution: Best distribution name for non-repairable
        
    Returns:
        tuple: (distribution_name, param1, param2)
    """
    if is_repairable:
        # For repairable equipment, method is in parameters["method"]
        method = parameters.get("method", "NHPP")
        
        if method == "NHPP":
            beta = parameters.get("beta", 1.0)
            eta = parameters.get("eta", 100000)
            return "NHPPTTFF", beta, eta
        elif method == "Weibull-2P":
            beta = parameters.get("beta", 1.0)
            eta = parameters.get("eta", 100000)
            return "Weibull2", beta, eta
        elif method == "Weibull-3P":
            beta = parameters.get("beta", 1.0)
            eta = parameters.get("eta", 100000)
            return "Weibull3", beta, eta
        elif method == "Exponential-2P":
            lambda_val = parameters.get("lambda_value", 0.00001)
            gamma = parameters.get("gamma", 0)
            return "Exponential2", lambda_val, gamma
        else:
            # Default to NHPP
            return "NHPPTTFF", 1.0, 100000
    else:
        # For non-repairable equipment, use best_distribution column
        distribution = best_distribution or "Normal"
        
        if distribution == "Lognormal":
            mu = parameters.get("mu", 0)
            sigma = parameters.get("sigma", 1)
            return "Lognormal", mu, sigma
        elif distribution == "Normal":
            mu = parameters.get("mu", 0)
            sigma = parameters.get("sigma", 1)
            return "Normal", mu, sigma
        elif distribution == "Weibull":
            # Assuming Weibull parameters are available
            beta = parameters.get("beta", 1.0)
            eta = parameters.get("eta", 100000)
            return "Weibull2", beta, eta
        else:
            # Default to Normal
            mu = parameters.get("mu", 0)
            sigma = parameters.get("sigma", 1)
            return "Normal", mu, sigma
    
async def update_oh_interval_offset(*, aeros_db_session, overhaul_offset, overhaul_interval, project_name):
    query = text("""
        UPDATE public."RegularNodes" rn
        SET "OHInterval" = :new_oh_interval,
            "OHOffset" = :new_oh_offset
        FROM public."Schematics" s
        JOIN public."Projects" p
        ON s."ProjectId" = p."ProjectId"
        WHERE rn."SchematicId" = s."SchematicId"
        AND p."ProjectName" = :project_name
    """)

    await aeros_db_session.execute(query, {
        "new_oh_interval": overhaul_interval,
        "new_oh_offset": overhaul_offset,
        "project_name": project_name
    })
    
    await aeros_db_session.commit()


async def update_equipment_for_simulation(
    *,
    db_session: DbSession,
    aeros_db_session: CollectorDbSession,
    project_name: str,
    simulation_duration: int,
    overhaul_duration,
    overhaul_interval,
    offset,
    schematic_name: str,
    custom_input: Optional[dict] = None
):
    log.info("Updating equipment for simulation")

    aeros_schematic = await get_aeros_schematic_by_name(
        db_session=db_session,
        schematic_name=schematic_name
    )

    equipments = (
        Select(AerosEquipment)
        .where(AerosEquipment.location_tag.isnot(None))
        .join(AerosEquipment.aeros_node)
        .filter(
            AerosNode.aeros_schematic_id == aeros_schematic.id,
            AerosNode.node_type == "RegularNode"
        )
    )

    rbd_group = Select(AerosEquipmentGroup)

    equipment_nodes = await db_session.execute(equipments)
    rbd_group_nodes = await db_session.execute(rbd_group)

    group_nodes = [g.group_name for g in rbd_group_nodes.scalars().all()]
    reg_nodes = [n.location_tag for n in equipment_nodes.scalars().all()]

    reg_nodes.extend(group_nodes)

    nodes_data = await get_aeros_equipment_by_location_tag(
        location_tag=reg_nodes,
        project_name=project_name
    )

    def load_json_file(filename: str):
        with open(filename, "r", encoding="utf-8") as f:
            return json.load(f)

    non_repairable = await db_session.execute(
        Select(ReliabilityPredictNonRepairable)
    )
    non_repairable_location_tags = [
        t.location_tag for t in non_repairable.scalars().all()
    ]

    reliability_data = await get_asset_batch(
        reg_nodes,
        non_repairable_location_tags,
        db_session
    )

    # Normalize custom input keys once
    custom_map = {}
    if custom_input:
        custom_map = {
            k.strip().upper(): v
            for k, v in custom_input.items()
        }
        log.info("Custom input provided for: %s", list(custom_map.keys()))

    reqNodeInputs = []
    results = {}

    await update_oh_interval_offset(
        aeros_db_session=aeros_db_session,
        project_name=project_name,
        overhaul_interval=overhaul_interval,
        overhaul_offset=offset
    )

    for eq in nodes_data:
        try:
            eq_name_raw = eq["equipmentName"]
            eq_name = eq_name_raw.strip().upper()

            log.debug("Processing equipment: %s", eq_name_raw)

            # ---- CUSTOM INPUT HAS PRIORITY ----
            if eq_name in custom_map:
                custom_param = custom_map[eq_name]

                mttr = custom_param.get("mttr")
                failure_rate = custom_param.get("failure_rate")

                if mttr is None or failure_rate is None:
                    log.warning(
                        "Custom input incomplete for %s, skipping",
                        eq_name_raw
                    )
                    continue

                if failure_rate == 0:
                    MTBF = 1000000
                else:
                    MTBF = 1e6 / float(failure_rate)

                eq["cmDisP1"] = mttr
                eq["relDisType"] = custom_param.get("distribution", "Fixed")

                if eq["relDisType"] == "Fixed":
                    eq["relDisP1"] = MTBF
                    eq["relDisP2"] = 0
                else:
                    eq["relDisP1"] = 1
                    eq["relDisP2"] = MTBF

                eq["ohDisP1"] = overhaul_duration
                eq["ohDisUnitCode"] = "UHour"

                reqNodeInputs.append(eq)
                results[eq_name_raw] = {
                    "mttr": mttr,
                    "distribution": eq["relDisType"],
                    "beta": eq["relDisP1"],
                    "eta": eq["relDisP2"],
                    "parameters": {},
                    "oh_duration": overhaul_duration
                }

                continue  # skip default logic

            # ---- PREVENT DUPLICATE DEFAULT INSERT ----
            if eq_name_raw in results:
                continue

            reliability = reliability_data.get(eq_name_raw, {})

            eq["cmDisType"] = reliability.get("cmDisType", "Fixed")
            eq["cmDisP1"] = reliability.get("cmDisP1", 0)
            eq["cmDisP2"] = reliability.get("cmDisP2", 0)
            eq["relDisType"] = reliability.get("relDisType", "Fixed")
            eq["relDisP1"] = reliability.get("relDisP1", 0)
            eq["relDisP2"] = reliability.get("relDisP2", 0)
            eq["ohDisP1"] = overhaul_duration
            eq["ohDisUnitCode"] = "UHour"

            reqNodeInputs.append(eq)
            results[eq_name_raw] = {
                "mttr": eq["cmDisP1"],
                "distribution": eq["relDisType"],
                "beta": eq["relDisP1"],
                "eta": eq["relDisP2"],
                "parameters": eq.get("parameters", {}),
                "oh_duration": overhaul_duration
            }

        except Exception as e:
            log.exception(
                "Error processing equipment %s", eq.get("equipmentName")
            )
            reqNodeInputs.append(eq)

    await update_node(
        db_session=db_session,
        equipment_nodes=reqNodeInputs,
        project_name=project_name
    )

    return results

# Optimized individual fetch functions
async def get_equipment_mttr(*, location_tag: str, client: httpx.AsyncClient) -> float:
    """
    Get MTTR for a single equipment using provided client
    """
    mttr_url = f"{RELIABILITY_SERVICE_API}/asset/mttr/{location_tag}"
    try:
        response = await client.get(mttr_url)
        if response.status_code == 200:
            mttr_data = response.json()
            return mttr_data.get("data", {}).get("hours", 0)
        else:
            log.warning(f"MTTR API returned status {response.status_code} for {location_tag}")
            return 0
    except Exception as e:
        log.error(f"Error getting MTTR for {location_tag}: {str(e)}")
        return 0

async def get_equipment_reliability_parameter(*, location_tag: str, client: httpx.AsyncClient):
    """
    Get reliability parameters for a single equipment using provided client
    """
    reliability_url = f"{RELIABILITY_SERVICE_API}/reliability/{location_tag}/current"
    try:
        response = await client.get(reliability_url)
        if response.status_code == 200:
            reliability_data = response.json()
            data = reliability_data.get("data", {})
            return {
                "distribution": data.get("distribution", ""),
                "beta": data.get("parameters", 0).get("beta", 0),
                "eta": data.get("parameters", 0).get("eta", 0)
            }
        else:
            log.warning(f"Reliability API returned status {response.status_code} for {location_tag}")
            return {"distribution": "", "beta": 0, "eta": 0}
    except Exception as e:
        log.error(f"Error getting reliability parameters for {location_tag}: {str(e)}")
        return {"distribution": "", "beta": 0, "eta": 0}
