# llm_integration/task_extraction.py
from openai import OpenAI
import re
import json

def extract_json_from_raw_response(raw_response):
    """
    Extract the JSON part from the raw response string.
    
    Args:
        raw_response (str): The raw response from the LLM containing JSON and additional text.

    Returns:
        dict: Parsed JSON object.
    """
    # Use regex to extract the JSON block between ```json and ```
    match = re.search(r"```json(.*?)```", raw_response, re.DOTALL)
    if match:
        json_string = match.group(1).strip()  # Extract the matched JSON part
        try:
            json_data = json.loads(json_string)  # Parse the JSON string into a Python dictionary
            return json_data
        except json.JSONDecodeError as e:
            print(f"Error decoding JSON: {e}")
            return None
    else:
        print("No valid JSON block found in the response.")
        return None
    

def compare_task_data(old_task_data, new_task_data):
    """
    Send old and new task data to the LLM for comparison.

    Args:
        old_task_data (dict): JSON data for the older tasks.
        new_task_data (dict): JSON data for the newer tasks.

    Returns:
        dict: Consolidated JSON with updates and new tasks.
    """
    # Prepare the prompt
    prompt = f"""
      Given the following two sets of task JSON data, compare them and:
      1. Identify projects and tasks present in the second JSON but not in the first. 
        - If two projects have different names but are contextually similar (e.g., due to spelling differences or tasks), treat them as the same project and merge their tasks.

      2. For tasks that exist in both JSONs within the same project:
        - Compare the following fields:
            - "description"
            - "priority"
            - "assigned_to"
            - "current_status"
        - If any changes are detected in these fields, update the task details in the output.

      3. If a project or task in the second JSON contains new tasks or subtasks not present in the first JSON:
        - Add those tasks or subtasks to the corresponding project in the output.

      4. Ensure the final JSON structure meets the following conditions:
        - Each project appears only once in the JSON.
        - All tasks are uniquely represented under their respective projects.
        - Updates to tasks (e.g., changes in "priority", "assigned_to", or "current_status") are applied.
        - Tasks or subtasks are not duplicated across the output.

      FIRST TASK DATA:
      '''
      {old_task_data}
      '''
      SECOND TASK DATA:
      '''
      {new_task_data}
      '''
      Expected Output:
      A single consolidated JSON structure where:
      - Projects are uniquely represented and merged based on contextual similarity.
      - Each project contains all relevant tasks, including updates and newly added ones.
      - All tasks follow this structure:

      Return a single consolidated JSON structure with:
      {{
        "project_name_1": {{
          "Task-1": {{
            "description": "Brief description of the task",
            "priority": "high/medium/low",
            "assigned_to": "Person responsible",
            "current_status": "Status of the task (e.g., completed, in progress, pending)"
          }},
          "Task-2": {{
            "description": "Brief description of the task",
            "priority": "high/medium/low",
            "assigned_to": "Person responsible",
            "current_status": "Status of the task (e.g., completed, in progress, pending)"
          }}
        }},
        "project_name_2": {{
          "Task-1": {{
            "description": "Brief description of the task",
            "priority": "high/medium/low",
            "assigned_to": "Person responsible",
            "current_status": "Status of the task (e.g., completed, in progress, pending)"
          }}
        }}
      }}
      """

    client = OpenAI(api_key='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')
    
    stream = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content":prompt}],
        # stream=True,
    )
    raw_response = stream.choices[0].message.content
    final_response= extract_json_from_raw_response(raw_response)
    return final_response