import json
import streamlit as st
import re

# Function to construct the prompt
def construct_prompt(prompt_template: str, user_input: str) -> str:
    return prompt_template.format(user_input=user_input)


def parse_json_content(cleaned_content: str):
    """
    Parses the cleaned content to extract valid JSON data.

    Args:
        cleaned_content (str): The raw content containing JSON data.

    Returns:
        dict or list: The parsed JSON object.
    """
    import re

    # Step 1: Strip unwanted characters and clean the content
    cleaned_content = cleaned_content.strip()

    # Step 2: Use regex to extract only the valid JSON block (e.g., starts with [ or {)
    json_match = re.search(r"(\{.*\}|\[.*\])", cleaned_content, re.DOTALL)
    
    if not json_match:
        raise ValueError("No valid JSON found in the content.")

    # Step 3: Extract and parse the valid JSON
    valid_json = json_match.group(0)  # Extract matched JSON block
    try:
        extracted_data = json.loads(valid_json)
    except json.JSONDecodeError as e:
        raise ValueError(f"Failed to decode JSON. Error: {e}\nContent:\n{valid_json}")

    return extracted_data

# Function to extract and parse JSON response
def extract_llm_response(response):
    """
    Extracts and parses the JSON response from the API.

    Args:
        response (dict): The API response containing a hint and proposals.

    Returns:
        tuple: A tuple containing the objective (str), key results (list), and hint (str).
    """

    raw_message_content = response["choices"][0]["message"]["content"]

    # Clean and parse the JSON content
    cleaned_content = raw_message_content.replace("`", "").split("json")[-1]

    parsed_data = parse_json_content(cleaned_content=cleaned_content)

    hint = parsed_data.get("hint", "")

    proposals = parsed_data.get("proposal", [])

    if proposals:
        # Extract the first proposal's objective and key results
        first_proposal = proposals[0]  # Get the first proposal (assuming it's a list)
        objective = first_proposal.get("objective", "")
        key_results = first_proposal.get("key_results", [])
    else:
        objective = ""
        key_results = []

    return objective, key_results, hint