import re

with open('src/calculation_time_constrains/service.py', 'r') as f:
    content = f.read()

# Replacement 1: run_simulation_with_spareparts
replacement_1 = """    finally:
        await optimum_oh_model._close_session()

    # Re-fetch calculation_data with equipment_results to ensure they are loaded
    from sqlalchemy import select
    from sqlalchemy.orm import selectinload
    calculation_query = await db_session.execute(
        select(CalculationData)
        .options(selectinload(CalculationData.equipment_results), selectinload(CalculationData.parameter))
        .where(CalculationData.id == calculation.id)
    )
    scope_calculation = calculation_query.scalar_one_or_none()

    data_num = scope_calculation.max_interval
    all_equipment = scope_calculation.equipment_results
    included_equipment = [eq for eq in all_equipment if eq.is_included]
    
    calculation_results = []
    fleet_statistics = {
        'total_equipment': len(all_equipment),
        'included_equipment': len(included_equipment),
        'excluded_equipment': len(all_equipment) - len(included_equipment),
        'equipment_with_sparepart_constraints': 0,
        'total_procurement_items': 0,
        'critical_procurement_items': 0,
        'total_spareparts': 745
    }

    avg_failure_cost = sum((eq.material_cost or 0) + (3 * 111000 * 3) for eq in all_equipment) / len(all_equipment) if all_equipment else 0

    rbd_marginal_fails = [0] * data_num
    try:
        if scope_calculation.rbd_simulation_id:
            from src.config import RBD_SERVICE_API
            import httpx
            plant_result_url = f"{RBD_SERVICE_API}/aeros/simulation/result/calc/{scope_calculation.rbd_simulation_id}/plant"
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.get(
                    plant_result_url,
                    headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
                )
                if response.status_code == 200:
                    plant_data = response.json().get("data", {})
                    timestamp_outs = plant_data.get("timestamp_outs", [])
                    if timestamp_outs:
                        from src.calculation_time_constrains.utils import create_time_series_data, calculate_failures_per_month
                        hourly_data = create_time_series_data(timestamp_outs, max_hours=data_num * 720)
                        cumulative_rbd_fails = calculate_failures_per_month(hourly_data)
                        rbd_fails_map = {m['month']: m['failures'] for m in cumulative_rbd_fails}
                        prev_fail = 0
                        for m in range(1, data_num + 1):
                            curr_fail = rbd_fails_map.get(m, prev_fail)
                            rbd_marginal_fails[m-1] = curr_fail - prev_fail
                            prev_fail = curr_fail
    except Exception as e:
        import logging
        logger = logging.getLogger(__name__)
        logger.warning(f"Failed to fetch plant simulation: {e}")

    cumulative_plant_failures = 0
    import numpy as np
    from .schema import CalculationResultsRead
    for month_index in range(data_num):
        historical_marginal_fail = 0
        for eq in all_equipment:
            if eq.is_actual and month_index < len(eq.is_actual) and eq.is_actual[month_index]:
                curr_fail = eq.daily_failures[month_index] if month_index < len(eq.daily_failures) else 0
                prev_fail = eq.daily_failures[month_index-1] if month_index > 0 and (month_index - 1) < len(eq.daily_failures) else 0
                historical_marginal_fail += max(0, curr_fail - prev_fail)

        marginal_fail = rbd_marginal_fails[month_index] + historical_marginal_fail
        cumulative_plant_failures += marginal_fail

        month_result = {
            "overhaul_cost": 0.0,
            "corrective_cost": 0.0,
            "procurement_cost": 0.0,
            "num_failures": cumulative_plant_failures,
            "day": month_index + 1,
            "month": month_index + 1,
            "procurement_details": {},
            "sparepart_summary": {
                "total_procurement_cost": 0.0,
                "equipment_requiring_procurement": 0,
                "critical_shortages": 0,
                "existing_orders_value": 0.0,
                "new_orders_required": 0,
                "urgent_procurements": 0
            }
        }

        equipment_requiring_procurement = 0
        total_existing_orders_value = 0.0
        total_new_orders_value = 0.0
        critical_shortages = 0
        urgent_procurements = 0

        for eq in all_equipment:
            if month_index >= len(eq.procurement_details):
                continue
            procurement_detail = eq.procurement_details[month_index]
            if (procurement_detail and isinstance(procurement_detail, dict) and procurement_detail.get("procurement_needed")):
                equipment_requiring_procurement += 1
                pr_po_summary = procurement_detail.get("pr_po_summary", {})
                existing_orders_value = pr_po_summary.get("total_existing_value", 0)
                total_existing_orders_value += existing_orders_value
                new_orders_value = pr_po_summary.get("total_new_orders_value", 0)
                total_new_orders_value += new_orders_value
                critical_missing = procurement_detail.get("critical_missing_parts", 0)
                if critical_missing > 0:
                    critical_shortages += 1
                recommendations = procurement_detail.get("recommendations", [])
                urgent_items = [r for r in recommendations if r.get("priority") == "CRITICAL"]
                if urgent_items:
                    urgent_procurements += 1
                is_included_eq = False if eq.is_initial else eq.is_included
                month_result["procurement_details"][eq.location_tag] = {
                    "is_included": is_included_eq,
                    "location_tag": eq.location_tag,
                    "details": procurement_detail.get("procurement_needed", []),
                    "detailed_message": procurement_detail.get("detailed_message", ""),
                    "pr_po_summary": pr_po_summary,
                    "recommendations": recommendations,
                    "sparepart_available": procurement_detail.get("sparepart_available", True),
                    "can_proceed": procurement_detail.get("can_proceed_with_delays", True),
                    "critical_missing_parts": critical_missing,
                    "existing_orders_value": existing_orders_value,
                    "new_orders_value": new_orders_value
                }
            if eq.is_included:
                if (month_index < len(eq.overhaul_costs) and month_index < len(eq.procurement_costs)):
                    month_result["overhaul_cost"] += float(eq.overhaul_costs[month_index])
                    month_result["procurement_cost"] += float(eq.procurement_costs[month_index])

        month_result["corrective_cost"] = (cumulative_plant_failures * avg_failure_cost) / (month_index + 1)
        month_result["sparepart_summary"].update({
            "total_procurement_cost": month_result["procurement_cost"],
            "equipment_requiring_procurement": equipment_requiring_procurement,
            "critical_shortages": critical_shortages,
            "existing_orders_value": total_existing_orders_value,
            "new_orders_required": len([eq for eq in all_equipment if month_index < len(eq.procurement_details) and eq.procurement_details[month_index] and eq.procurement_details[month_index].get("procurement_needed")]),
            "urgent_procurements": urgent_procurements
        })
        month_result["total_cost"] = month_result["corrective_cost"] + month_result["overhaul_cost"] + month_result["procurement_cost"]
        calculation_results.append(month_result)

    optimum_day = np.argmin([month["total_cost"] for month in calculation_results])
    scope_calculation.optimum_oh_day = int(optimum_day)
    
    fleet_statistics['equipment_with_sparepart_constraints'] = len([
        eq for eq in all_equipment if any(detail and detail.get("procurement_needed") for detail in eq.procurement_details if detail)
    ])
    fleet_statistics['total_procurement_items'] = sum([
        len(detail.get("procurement_needed", [])) for eq in all_equipment for detail in eq.procurement_details if detail and isinstance(detail, dict)
    ])

    analysis_metadata = {
        "planned_month": (scope.start_date.year - prev_oh_scope.end_date.year) * 12 + (scope.start_date.month - prev_oh_scope.end_date.month) if prev_oh_scope and scope else 0,
        "total_fleet_procurement_cost": sum([eq.procurement_costs[int(scope_calculation.optimum_oh_day)] for eq in all_equipment if eq.procurement_costs]),
        "last_oh_date": prev_oh_scope.end_date.isoformat() if prev_oh_scope else None,
        "next_oh_date": scope.start_date.isoformat() if scope else None,
        "optimal_stat": None
    }
    
    calc_results_read = [CalculationResultsRead(**r) for r in calculation_results]
    optimal_analysis = _analyze_optimal_timing(
        calc_results_read, scope_calculation.optimum_oh_day, prev_oh_scope, scope
    )

    scope_calculation.plant_results = calculation_results
    scope_calculation.fleet_statistics = fleet_statistics
    scope_calculation.analysis_metadata = analysis_metadata
    scope_calculation.optimum_analysis = optimal_analysis

    await db_session.commit()

    return {
        "id": calculation.id,
        "optimum": optimal_analysis
    }"""

pattern_1 = re.compile(r"    finally:\n        await optimum_oh_model\._close_session\(\).*?return \{\n        \"id\": calculation_data\.id,\n        \"optimum\": stats\n    \}", re.DOTALL)
if pattern_1.search(content):
    content = pattern_1.sub(replacement_1, content)
else:
    print("Could not find Replacement 1 target")

# Replacement 2: get_calculation_result
replacement_2 = """async def get_calculation_result(db_session: DbSession, calculation_id: str, token, include_risk_cost):
    \"\"\"
    Get calculation results from DB, returning pre-calculated plant and equipment results.
    \"\"\"
    try:
        # Get calculation data with equipment results
        calculation_query = await db_session.execute(
            select(CalculationData)
            .options(selectinload(CalculationData.equipment_results))
            .where(CalculationData.id == calculation_id)
        )
        scope_calculation = calculation_query.scalar_one_or_none()
        
        if not scope_calculation:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"Calculation with id {calculation_id} does not exist.",
            )

        scope_overhaul = await get_scope(
            db_session=db_session, 
            overhaul_session_id=scope_calculation.overhaul_session_id
        )
        if not scope_overhaul:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=f"Overhaul scope for session {scope_calculation.overhaul_session_id} does not exist.",
            )

        # Parse pre-calculated plant results
        plant_results_raw = scope_calculation.plant_results or []
        calculation_results = [CalculationResultsRead(**r) for r in plant_results_raw]

        # Return comprehensive result
        return CalculationTimeConstrainsRead(
            id=scope_calculation.id,
            reference=scope_calculation.overhaul_session_id,
            scope=scope_overhaul.maintenance_type.name,
            results=calculation_results,
            optimum_oh=scope_calculation.optimum_oh_day,
            optimum_oh_month=scope_calculation.optimum_oh_day + 1 if scope_calculation.optimum_oh_day is not None else None,
            equipment_results=scope_calculation.equipment_results,
            fleet_statistics=scope_calculation.fleet_statistics or {},
            optimal_analysis=scope_calculation.optimum_analysis or {},
            analysis_metadata=scope_calculation.analysis_metadata or {}
        )

    except HTTPException:
        raise
    except Exception as e:
        import logging
        logger = logging.getLogger(__name__)
        logger.error(f"Error in get_calculation_result for calculation_id {calculation_id}: {str(e)}")
        
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Internal error processing calculation results: {str(e)}",
        )"""

pattern_2 = re.compile(r"async def get_calculation_result\(db_session: DbSession, calculation_id: str, token, include_risk_cost\):.*?raise HTTPException\(\n            status_code=status\.HTTP_500_INTERNAL_SERVER_ERROR,\n            detail=f\"Internal error processing calculation results: \{str\(e\)\}\",\n        \)", re.DOTALL)
if pattern_2.search(content):
    content = pattern_2.sub(replacement_2, content)
else:
    print("Could not find Replacement 2 target")

with open('src/calculation_time_constrains/service.py', 'w') as f:
    f.write(content)

print("Patch applied.")