class LLMTM():
    def __init__(self, LLM_model):
        self.LLM_model = LLM_model
    
    def task_knowledge_refinement(self, task_description, web_search=False):
        # 1. Create the prompt for knowledge refinement 
        prompt = f"""
You are a world-class subject matter expert with deep knowledge related to the following task.
Your goal is to provide essential background knowledge, expert insights, and context that will help someone better understand and execute this task.

Please analyze the task description below and generate a concise, informative expansion of the key concepts.
Do NOT repeat the original task description in your output. Focus only on adding new, relevant expert knowledge or detailed explanations for the key concepts.

**Original Task Description:**
---
{task_description}
---

**Your Expert Knowledge Expansion:**
"""

        # 2. Use LLM to generate output
        additional_knowledge = self.LLM_model.LLM_response(prompt)
        extracted_knowledge = additional_knowledge.strip()

        # 3. Combine the original task description and the additional knowledge
        # We format the output clearly using Markdown for better readability.
        refined_task_description = (
            "\n\n## Task Description\n\n"
            f"{task_description}"
            "\n\n"
            "## Expert Knowledge & Context\n\n"
            f"{extracted_knowledge}"
        )

        # 4. Return the refined task description
        return refined_task_description

    def get_plan(self, task_description, num_plans):
        print(f"INFO: Generating {num_plans} initial plans...")
        plans = {f"plan_{i}": f"Initial plan content {i}" for i in range(num_plans)}
        return plans, plans # Returning raw plans as well
    
    def get_plan(self, task_description, num_plans=1):
        if num_plans < 1:
            return []

        # 1. Generate the planning prompt from the helper function
        plan_prompt = self.get_planning_prompt(task_description)

        # 2. Create a list of prompts to generate multiple plans in parallel
        prompts_list = [plan_prompt] * num_plans

        # 3. Call the asynchronous LLM function to get plans efficiently
        results, all_successful = self.LLM_model.LLM_response_async(prompts=prompts_list)

        if not all_successful:
            print("Warning: some plan generation requests failed.")

        # 4. Collect and return the successfully generated plans
        final_plans = [response for success, response, _ in results if success]

        return final_plans
    
    def summarize_plan(self, plans, single_plan=True):
        if isinstance(plans, str):
            plans = [plans]

        if single_plan and len(plans) > 1:
            plans = [plans[0]]

        # 1. Generate the detailed prompt using the helper function
        prompt = self.plan_summarize_prompt(plans)

        # 2. Call the LLM to get the summarization.
        try:
            summarized_plan = self.LLM_model.LLM_response(prompt)
            return summarized_plan.strip()
        except Exception as e:
            # Handle potential exceptions during the LLM call
            print(f"An error occurred during LLM call for plan summarization: {e}")
            return ""

    def plan_refinement_loop(self, plan, max_refine_rounds=3):
        current_plan = ""

        # --- Step 1: Handle Input Type and Select a Single Plan ---
        if isinstance(plan, list):
            if not plan:
                print("Error: The provided plan list is empty.")
                return ""
            
            # If only one plan is in the list, select it automatically
            if len(plan) == 1:
                current_plan = plan[0]
            else:
                # Translate all plans in the list for user review
                translation_prompts = [self.get_translation_prompt(p) for p in plan]
                results, _ = self.LLM_model.LLM_response_async(prompts=translation_prompts)

                # Correlate successful translations with their original technical plans
                valid_options = [(plan[i], resp) for i, (succ, resp, _) in enumerate(results) if succ]

                if not valid_options:
                    print("Error: Failed to translate any of the provided plans.")
                    return ""

                # Present the translated plans to the user for selection
                print("\nPlease select one of the following plans to proceed with:")
                for i, (_, readable_plan) in enumerate(valid_options):
                    print(f"\n--- [Option {i + 1}] ---")
                    print(readable_plan)
                    print("--------------------")

                # Get and validate user's choice
                choice = -1
                while True:
                    try:
                        user_input = input(f"\nEnter the number of the plan you want to refine (1-{len(valid_options)}): ")
                        choice = int(user_input)
                        if 1 <= choice <= len(valid_options):
                            break
                        else:
                            print(f"Invalid selection. Please enter a number between 1 and {len(valid_options)}.")
                    except ValueError:
                        print("Invalid input. Please enter a number.")
                
                # Set the chosen technical plan as the one to be refined
                current_plan = valid_options[choice - 1][0]

        elif isinstance(plan, str):
            current_plan = plan
        else:
            raise TypeError("The 'plan' argument must be a string or a list of strings.")

        # --- Step 2: Begin the Refinement Loop on the Selected Plan ---
        for i in range(max_refine_rounds):
            print(f"\n--- [Refinement Round {i + 1}/{max_refine_rounds}] ---")

            # Translate the current technical plan into a readable format
            translation_prompt = self.get_translation_prompt(current_plan)
            readable_plan = self.LLM_model.LLM_response(translation_prompt)

            # Show the readable plan to the user
            print("\nHere is the current plan for your review:")
            print("-----------------------------------------")
            print(readable_plan)
            print("-----------------------------------------")

            # Ask if the user wants to refine it further
            user_input = input("Are you satisfied with this plan, or would you like to refine it? (yes/no refine): ").lower().strip()

            if user_input.startswith('y'):
                print("\nPlan approved. Finalizing the process.")
                return current_plan # Return the approved technical plan

            # Get user feedback for refinement
            feedback = input("\nPlease provide your feedback for refinement:\n> ")

            # Refine the plan based on the feedback
            print("\nRefining the plan based on your feedback...")
            refinement_prompt = self.get_refinement_prompt(
                original_plan=current_plan,
                user_feedback=feedback
            )
            refined_plan = self.LLM_model.LLM_response(refinement_prompt)

            # The refined plan becomes the new "current" plan for the next loop
            current_plan = refined_plan

        print(f"\nMaximum refinement rounds ({max_refine_rounds}) reached. Using the last generated plan.")
        return current_plan

    def get_planning_prompt(self, task_description):
        prompt = f"""
You are an expert software architect and planning agent. Your primary role is to analyze a user's coding request and generate a comprehensive, step-by-step implementation plan. This plan will be used as a blueprint by a separate code generation agent to write the actual code.

Your plan must be clear, logical, and detailed enough for the code agent to understand and implement without further clarification.

**User's Task Description:**
"{task_description}"

---

Please generate the implementation plan by structuring your response into the following three sections:

### 1. Overall Task Definition
Summarize the core goal of the task. Clearly and precisely define the expected **input format** for the entire program and the **final output format** that the program must produce. This ensures the final solution can be tested correctly.

### 2. Component Breakdown
Break the problem down into logical, self-contained components (e.g., functions, classes). For each component, you must provide:
- **Name:** A descriptive name for the component (e.g., `parse_user_data`).
- **Purpose:** A brief explanation of what this component does.
- **Inputs:** A clear description of the arguments it takes, including their expected data types.
- **Outputs:** A clear description of what it returns, including its data type.

### 3. Main Function (Dataflow)
Describe the logic and workflow for the `main` function. This function will serve as the entry point for the entire script and will be used for testing and evaluation. Your description should detail the dataflow:
- How the initial input is received by the `main` function.
- The sequence in which the components defined in Part 2 are called.
- How the output of one component becomes the input for another.
- How the final result is assembled and returned, matching the format specified in the "Overall Task Definition"."""
        return prompt

    def get_refinement_prompt(self, original_plan, user_feedback):
        prompt = f"""
You are an expert software architect. Your task is to refine an existing implementation plan based on user feedback. You must generate a new, complete plan that incorporates the requested changes while maintaining the original structure.

**Original Plan:**
---
{original_plan}
---

**User's Feedback for Refinement:**
---
{user_feedback}
---

Please generate the **full, updated plan** that addresses the user's feedback. The new plan must be a complete replacement for the original and follow the same three-part format:

### 1. Overall Task Definition
(Update the input/output definitions if the feedback requires it.)

### 2. Component Breakdown
(Modify, add, or remove components based on the feedback.)

### 3. Main Function (Dataflow)
(Update the workflow to reflect the changes in the components and logic.)
"""
        return prompt

    def get_translation_prompt(self, technical_plan):
        prompt = f"""
You are an expert communicator specializing in translating technical software plans into plain language for non-technical clients.

The user is an expert in their subject matter but has no coding experience. Your goal is to explain the proposed software's workflow clearly and concisely.

**Instructions:**
1.  Read the provided technical plan carefully.
2.  Do NOT use programming jargon (e.g., 'function', 'class', 'API', 'data type', 'return value').
3.  Instead, use analogies: a 'component' or 'function' can be a 'processing step' or a 'tool'; 'input' is 'information needed'; 'output' is the 'result produced'.
4.  Structure your explanation in a clear, step-by-step format that is easy to follow.

**Technical Plan to Translate:**
---
{technical_plan}
---

**Generate the user-friendly explanation below, focusing on:**
1.  **Overall Goal:** A brief, one-sentence summary of what the program will accomplish.
2.  **Required Information:** What information the program needs to start its work.
3.  **The Process (Step-by-Step):** Describe each component from the plan as a distinct step in a workflow. For each step, explain its purpose, what information it uses, and what result it produces.
4.  **Final Result:** Describe the final output the user will receive from the program.
"""
        return prompt

    def generate_new_plan(self, best_plan, task_desc, codes, tests):
        print(f"INFO: Generating new plans based on the current best: {best_plan}")
        return {f"{best_plan}_v2": "A new, improved plan version"}
    
    def plan_summarize_prompt(self, plans):
        # Ensure plans is a list for consistent processing
        if isinstance(plans, str):
            plans = [plans]

        # Format the input plans for inclusion in the prompt
        formatted_plans = ""
        if len(plans) == 1:
            formatted_plans = f"## Plan to be Summarized:\n\n{plans[0]}"
        else:
            # If multiple plans, number them for clarity
            for i, plan in enumerate(plans):
                formatted_plans += f"## Plan {i+1}:\n\n{plan}\n\n---\n\n"

        # The core prompt template instructing the LLM on its task
        prompt = f"""
You are an expert technical analyst. Your task is to read and synthesize the following software development plan(s). Your output will be a single, consolidated summary that will be used by another AI agent to generate comprehensive test cases.

The summary must be **neutral and objective**. If multiple plans are provided, combine their requirements without showing preference for any single implementation strategy. Focus strictly on the **"what" (requirements)**, not the "how" (implementation details).

From the provided plan(s), extract and structure your summary with the following sections:

1.  **Main Task:** A concise, one-sentence description of the overall goal.
2.  **Input Specification:** Detail all expected inputs. For each input, specify its name, data type (e.g., string, integer, list of dictionaries), format (e.g., JSON object, CSV string), and any known constraints or validation rules (e.g., "integer must be positive", "string cannot be empty").
3.  **Output Specification:** Detail the expected output. Specify its data type, format, and any conditions it must satisfy (e.g., "returns a dictionary with 'status' and 'data' keys", "returns -1 on error").
4.  **Core Functionalities & Logic:** List the key behaviors, logical steps, or calculations the code must perform. This helps identify different scenarios to test.
5.  **Constraints & Edge Cases:** Explicitly list any system limitations, assumptions, or specific edge cases mentioned in the plans (e.g., handling of null/empty inputs, maximum file size, error handling conditions).

Do **not** suggest how to write test cases. Only provide the factual specifications derived from the plan(s) below.

---
{formatted_plans}
---

Provide the structured summary now.
"""
        return prompt.strip()