import pandas as pd
import litellm
import dotenv
import os
import time
import json
import math
import numpy as np

# --- Configuration ---
MODEL = "gpt-4.1-mini"  # Make sure this model supports json_schema or structured output
RATE_LIMIT = 5000  # Requests per minute
CHUNK_SIZE = 300
SECONDS_PER_MINUTE = 60
FILENAME = (
    "tasks_with_estimates.csv"  # This CSV should contain the tasks to be processed
)

# --- Prompts and Schema ---
SYSTEM_PROMPT = """
You are an expert assistant evaluating the time to completion required for job tasks. Your goal is to estimate the time range needed for a skilled human to complete the following job task remotely, without supervision.

Provide a lower and upper bound estimate for the time to completion time. These bounds should capture the time within which approximately 80% of instances of performing this specific task are typically completed by a qualified individual.

Base your estimate on the provided task description, its associated activities, and the occupational context. Your estimate must be in one the allowed units: minute, hour, day, week, month, trimester, semester, year.
""".strip()

USER_MESSAGE_TEMPLATE = """
Please estimate the time range for the following remote task:

**Task Description:** {task}
**Relevant activies for the task:**
{dwas}

**Occupation Category:** {occupation_title}
**Occupation Description:** {occupation_description}

Consider the complexity and the typical steps involved.
""".strip()

ALLOWED_UNITS = [
    "minute",
    "hour",
    "day",
    "week",
    "month",
    "trimester",
    "semester",
    "year",
]

SCHEMA_FOR_VALIDATION = {
    "name": "estimate_time",
    "strict": True,  # Enforce schema adherence
    "schema": {
        "type": "object",
        "properties": {
            "lower_bound_estimate": {
                "type": "object",
                "properties": {
                    "quantity": {
                        "type": "number",
                        "description": "The numerical value for the lower bound of the estimate.",
                    },
                    "unit": {
                        "type": "string",
                        "enum": ALLOWED_UNITS,
                        "description": "The unit of time for the lower bound.",
                    },
                },
                "required": ["quantity", "unit"],
                "additionalProperties": False,
            },
            "upper_bound_estimate": {
                "type": "object",
                "properties": {
                    "quantity": {
                        "type": "number",
                        "description": "The numerical value for the upper bound of the estimate.",
                    },
                    "unit": {
                        "type": "string",
                        "enum": ALLOWED_UNITS,
                        "description": "The unit of time for the upper bound.",
                    },
                },
                "required": ["quantity", "unit"],
                "additionalProperties": False,
            },
        },
        "required": ["lower_bound_estimate", "upper_bound_estimate"],
        "additionalProperties": False,
    },
}


def save_dataframe(df_to_save, filename):

    """Saves the DataFrame to the specified CSV file using atomic write."""
    try:
        temp_filename = filename + ".tmp"
        df_to_save.to_csv(temp_filename, encoding="utf-8-sig", index=False)
        os.replace(temp_filename, filename)
    except Exception as e:
        print(f"--- Error saving DataFrame to {filename}: {e} ---")
        if os.path.exists(temp_filename):
            try:
                os.remove(temp_filename)
            except Exception as remove_err:
                print(
                    f"--- Error removing temporary save file {temp_filename}: {remove_err} ---"
                )

def create_task_estimates():
    try:
        # Read the CSV
        if os.path.exists(FILENAME):
            df = pd.read_csv(FILENAME, encoding="utf-8-sig")
            print(f"Successfully read {len(df)} rows from {FILENAME}.")

            estimate_columns_spec = {
                "lb_estimate_qty": float,
                "lb_estimate_unit": object,
                "ub_estimate_qty": float,
                "ub_estimate_unit": object,
            }
            save_needed = False

            for col_name, target_dtype in estimate_columns_spec.items():
                if col_name not in df.columns:
                    # Initialize with a type-compatible missing value
                    if target_dtype == float:
                        df[col_name] = np.nan
                    else:  # object
                        df[col_name] = pd.NA
                    df[col_name] = df[col_name].astype(target_dtype)  # Enforce dtype
                    print(f"Added '{col_name}' column as {df[col_name].dtype}.")
                    save_needed = True
                else:
                    # Column exists, ensure correct dtype
                    current_pd_dtype = df[col_name].dtype
                    expected_pd_dtype = pd.Series(dtype=target_dtype).dtype

                    if current_pd_dtype != expected_pd_dtype:
                        try:
                            if target_dtype == float:
                                df[col_name] = pd.to_numeric(df[col_name], errors="coerce")
                            else:  # object
                                df[col_name] = df[col_name].astype(object)
                            print(
                                f"Corrected dtype of '{col_name}' to {df[col_name].dtype}."
                            )
                            save_needed = True
                        except Exception as e:
                            print(
                                f"Warning: Could not convert column '{col_name}' to {target_dtype}: {e}. Current dtype: {current_pd_dtype}"
                            )

                # Standardize missing values (e.g., empty strings to NA/NaN)
                # Replace common missing placeholders with pd.NA first
                df[col_name].replace(["", None, ""], pd.NA, inplace=True)
                if target_dtype == float:
                    # For float columns, ensure they are numeric and use np.nan after replacement
                    df[col_name] = pd.to_numeric(df[col_name], errors="coerce")

            if save_needed:
                print(f"Saving {FILENAME} after adding/adjusting estimate columns.")
                save_dataframe(df, FILENAME)
        else:
            print(
                f"Error: {FILENAME} not found. Please ensure the file exists and contains task data."
            )
            exit()
    except FileNotFoundError:
        print(
            f"Error: {FILENAME} not found. Please ensure the file exists and contains task data."
        )
        exit()
    except Exception as e:
        print(f"Error reading or initializing {FILENAME}: {e}")
        exit()

    # --- Identify Rows to Process ---
    # We'll check for NaN in one of the primary quantity columns.
    unprocessed_mask = df["lb_estimate_qty"].isna()
    if unprocessed_mask.any():
        start_index = unprocessed_mask.idxmax()  # Finds the index of the first True value
        print(f"Resuming processing. First unprocessed row found at index {start_index}.")
        df_to_process = df.loc[unprocessed_mask].copy()
        original_indices = df_to_process.index  # Keep track of original indices
    else:
        print(
            "All rows seem to have estimates already (based on 'lb_estimate_qty'). Exiting."
        )
        exit()


    # --- Prepare messages for batch completion (only for rows needing processing) ---
    messages_list = []
    skipped_rows_indices = []
    valid_original_indices = []

    if not df_to_process.empty:
        required_cols = ["task", "occupation_title", "occupation_description", "dwas"]
        print(
            f"Preparing messages for up to {len(df_to_process)} rows starting from original index {original_indices[0] if len(original_indices) > 0 else 'N/A'}..."
        )
        print(f"Checking for required columns: {required_cols}")

        for index, row in df_to_process.iterrows():
            missing_or_empty = []
            for col in required_cols:
                if col not in row or pd.isna(row[col]) or str(row[col]).strip() == "":
                    missing_or_empty.append(col)

            if missing_or_empty:
                print(
                    f"Warning: Skipping row original index {index} due to missing/empty required data in columns: {', '.join(missing_or_empty)}."
                )
                skipped_rows_indices.append(index)
                continue

            try:
                user_message = USER_MESSAGE_TEMPLATE.format(
                    task=row["task"],
                    occupation_title=row["occupation_title"],
                    occupation_description=row["occupation_description"],
                    dwas=row["dwas"],
                )
            except KeyError as e:
                print(
                    f"Error: Skipping row original index {index} due to formatting error - missing key: {e}. Check USER_MESSAGE_TEMPLATE and CSV columns."
                )
                skipped_rows_indices.append(index)
                continue

            messages_for_row = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_message},
            ]
            messages_list.append(messages_for_row)
            valid_original_indices.append(index)  # This is the original DataFrame index

        print(
            f"Prepared {len(messages_list)} valid message sets for batch completion (skipped {len(skipped_rows_indices)} rows)."
        )
        if not messages_list:
            print("No valid rows found to process after checking required data. Exiting.")
            exit()
    else:
        print(
            "No rows found needing processing (df_to_process is empty)."
        )  # Should have been caught by earlier check
        exit()


    # --- Call batch_completion in chunks with rate limiting and periodic saving ---
    total_messages_to_send = len(messages_list)
    num_chunks = math.ceil(total_messages_to_send / CHUNK_SIZE)

    print(
        f"\nStarting batch completion for {total_messages_to_send} items in {num_chunks} chunks..."
    )

    overall_start_time = time.time()
    processed_count_total = 0

    for i in range(num_chunks):
        chunk_start_message_index = i * CHUNK_SIZE
        chunk_end_message_index = min((i + 1) * CHUNK_SIZE, total_messages_to_send)
        message_chunk = messages_list[chunk_start_message_index:chunk_end_message_index]
        # Get corresponding original DataFrame indices for this chunk
        chunk_original_indices = valid_original_indices[
            chunk_start_message_index:chunk_end_message_index
        ]

        if not message_chunk:
            continue

        min_idx_disp = min(chunk_original_indices) if chunk_original_indices else "N/A"
        max_idx_disp = max(chunk_original_indices) if chunk_original_indices else "N/A"
        print(
            f"\nProcessing chunk {i + 1}/{num_chunks} (Messages {chunk_start_message_index + 1}-{chunk_end_message_index} of this run)..."
            f" Corresponding to original indices: {min_idx_disp} - {max_idx_disp}"
        )
        chunk_start_time = time.time()
        responses = []
        try:
            print(f"Sending {len(message_chunk)} requests for chunk {i + 1}...")
            responses = litellm.batch_completion(
                model=MODEL,
                messages=message_chunk,
                response_format={
                    "type": "json_schema",
                    "json_schema": SCHEMA_FOR_VALIDATION,
                },
                num_retries=3,
                # request_timeout=60 # Optional: uncomment if needed
            )
            print(f"Chunk {i + 1} API call completed.")

        except Exception as e:
            print(f"Error during litellm.batch_completion for chunk {i + 1}: {e}")
            responses = [None] * len(
                message_chunk
            )  # Ensure responses list matches message_chunk length for processing loop

        # --- Process responses for the current chunk ---
        chunk_updates = {}  # To store {original_df_index: {qty/unit data}}
        successful_in_chunk = 0
        failed_in_chunk = 0

        if responses and len(responses) == len(message_chunk):
            for j, response in enumerate(responses):
                original_df_index = chunk_original_indices[j]

                # Initialize values for this item
                lb_qty_val, lb_unit_val, ub_qty_val, ub_unit_val = None, None, None, None
                content_str = None

                if response is None:
                    print(
                        f"Skipping processing for original index {original_df_index} due to API call failure for this item (response is None)."
                    )
                    failed_in_chunk += 1
                    continue

                try:
                    if (
                        response.choices
                        and response.choices[0].message
                        and response.choices[0].message.content
                    ):
                        content_str = response.choices[0].message.content
                        estimate_data = json.loads(content_str)  # Can raise JSONDecodeError

                        lower_bound_dict = estimate_data.get("lower_bound_estimate")
                        upper_bound_dict = estimate_data.get("upper_bound_estimate")

                        valid_response_structure = isinstance(
                            lower_bound_dict, dict
                        ) and isinstance(upper_bound_dict, dict)

                        if valid_response_structure:
                            lb_qty_raw = lower_bound_dict.get("quantity")
                            lb_unit_raw = lower_bound_dict.get("unit")
                            ub_qty_raw = upper_bound_dict.get("quantity")
                            ub_unit_raw = upper_bound_dict.get("unit")

                            is_valid_item = True
                            # Validate LB Qty
                            if (
                                not isinstance(lb_qty_raw, (int, float))
                                or math.isnan(float(lb_qty_raw))
                                or float(lb_qty_raw) < 0
                            ):
                                print(
                                    f"Warning: Invalid lb_quantity for original index {original_df_index}: {lb_qty_raw}"
                                )
                                is_valid_item = False
                            else:
                                lb_qty_val = float(lb_qty_raw)

                            # Validate UB Qty
                            if (
                                not isinstance(ub_qty_raw, (int, float))
                                or math.isnan(float(ub_qty_raw))
                                or float(ub_qty_raw) < 0
                            ):
                                print(
                                    f"Warning: Invalid ub_quantity for original index {original_df_index}: {ub_qty_raw}"
                                )
                                is_valid_item = False
                            else:
                                ub_qty_val = float(ub_qty_raw)

                            # Validate Units
                            if lb_unit_raw not in ALLOWED_UNITS:
                                print(
                                    f"Warning: Invalid lb_unit for original index {original_df_index}: '{lb_unit_raw}'"
                                )
                                is_valid_item = False
                            else:
                                lb_unit_val = lb_unit_raw

                            if ub_unit_raw not in ALLOWED_UNITS:
                                print(
                                    f"Warning: Invalid ub_unit for original index {original_df_index}: '{ub_unit_raw}'"
                                )
                                is_valid_item = False
                            else:
                                ub_unit_val = ub_unit_raw

                            if is_valid_item:
                                successful_in_chunk += 1
                                chunk_updates[original_df_index] = {
                                    "lb_estimate_qty": lb_qty_val,
                                    "lb_estimate_unit": lb_unit_val,
                                    "ub_estimate_qty": ub_qty_val,
                                    "ub_estimate_unit": ub_unit_val,
                                }
                            else:
                                failed_in_chunk += (
                                    1  # Values remain None if not fully valid
                                )
                        else:
                            print(
                                f"Warning: Missing or malformed estimate dicts in JSON for original index {original_df_index}. Content: '{content_str}'"
                            )
                            failed_in_chunk += 1
                    else:
                        finish_reason = (
                            response.choices[0].finish_reason
                            if (response.choices and response.choices[0].finish_reason)
                            else "unknown"
                        )
                        error_message = (
                            response.choices[0].message.content
                            if (
                                response.choices
                                and response.choices[0].message
                                and response.choices[0].message.content
                            )
                            else "No content in message."
                        )
                        print(
                            f"Warning: Received non-standard or empty response content for original index {original_df_index}. "
                            f"Finish Reason: '{finish_reason}'. Message: '{error_message}'. Raw Choices: {response.choices}"
                        )
                        failed_in_chunk += 1

                except json.JSONDecodeError:
                    print(
                        f"Warning: Could not decode JSON for original index {original_df_index}. Content received: '{content_str}'"
                    )
                    failed_in_chunk += 1
                except AttributeError as ae:
                    print(
                        f"Warning: Missing expected attribute processing response for original index {original_df_index}: {ae}. Response: {response}"
                    )
                    failed_in_chunk += 1
                except Exception as e:
                    print(
                        f"Warning: An unexpected error occurred processing response for original index {original_df_index}: {type(e).__name__} - {e}. Response: {response}"
                    )
                    failed_in_chunk += 1
        else:
            print(
                f"Warning: Mismatch between number of responses ({len(responses) if responses else 0}) "
                f"and messages sent ({len(message_chunk)}) for chunk {i + 1}, or no responses. Marking all as failed."
            )
            failed_in_chunk = len(
                message_chunk
            )  # All items in this chunk are considered failed if response array is problematic

        print(
            f"Chunk {i + 1} processing summary: Success={successful_in_chunk}, Failed/Skipped={failed_in_chunk}"
        )
        processed_count_total += successful_in_chunk

        # --- Update Main DataFrame and Save Periodically ---
        if chunk_updates:
            print(
                f"Updating main DataFrame with {len(chunk_updates)} new estimates for chunk {i + 1}..."
            )
            for idx, estimates in chunk_updates.items():
                if idx in df.index:
                    df.loc[idx, "lb_estimate_qty"] = estimates["lb_estimate_qty"]
                    df.loc[idx, "lb_estimate_unit"] = estimates["lb_estimate_unit"]
                    df.loc[idx, "ub_estimate_qty"] = estimates["ub_estimate_qty"]
                    df.loc[idx, "ub_estimate_unit"] = estimates["ub_estimate_unit"]

            print(f"Saving progress to {FILENAME}...")
            save_dataframe(df, FILENAME)
        else:
            print(f"No successful estimates obtained in chunk {i + 1} to save.")

        # --- Rate Limiting Pause ---
        chunk_end_time = time.time()
        chunk_duration = chunk_end_time - chunk_start_time
        print(f"Chunk {i + 1} took {chunk_duration:.2f} seconds.")

        if i < num_chunks - 1:  # No pause after the last chunk
            # Calculate ideal time per request based on rate limit
            time_per_request = SECONDS_PER_MINUTE / RATE_LIMIT if RATE_LIMIT > 0 else 0
            # Calculate minimum duration this chunk should have taken to respect rate limit
            min_chunk_duration_for_rate = len(message_chunk) * time_per_request
            # Calculate pause needed
            pause_needed = max(0, min_chunk_duration_for_rate - chunk_duration)

            if pause_needed > 0:
                print(
                    f"Pausing for {pause_needed:.2f} seconds to respect rate limit ({RATE_LIMIT}/min)..."
                )
                time.sleep(pause_needed)

    overall_end_time = time.time()
    total_duration_minutes = (overall_end_time - overall_start_time) / 60
    print(
        f"\nBatch completion finished."
        f" Processed {processed_count_total} new estimates in this run in {total_duration_minutes:.2f} minutes."
    )

    print(f"Performing final save to {FILENAME}...")
    save_dataframe(df, FILENAME)

    print("\nScript finished.")
