import ollama
import re
# def clean_json_response(raw_string):
#     """
#     Removes Markdown code blocks and extra whitespace from AI responses.
#     """
#     # Use regex to find content between ```json and ``` or just ```
#     match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', raw_string)
#     if match:
#         return match.group(1).strip()
#     return raw_string.strip()

import json
import re

def robust_json_parser(raw_output):
    # 1. Strip Markdown blocks
    clean = re.sub(r'```(?:json)?|```', '', raw_output).strip()
    
    # 2. Fix trailing commas before closing braces/brackets
    clean = re.sub(r',\s*([\]}])', r'\1', clean)
    
    try:
        return json.loads(clean)
    except json.JSONDecodeError:
        # 3. Last resort: try to find the first '{' and last '}'
        start_idx = clean.find('{')
        end_idx = clean.rfind('}')
        if start_idx != -1 and end_idx != -1:
            try:
                return json.loads(clean[start_idx:end_idx+1])
            except:
                pass
        raise

def get_model_reponse(prompt):
    response=ollama.chat(
        model='alibayram/smollm3:latest',
        messages=[{'role': 'user', 'content': prompt}],
        stream=False # Set to True for real-time streaming
    )
    # print(response['message']['content'])
    response=robust_json_parser(response['message']['content'])
    return response