import json
import traceback
from typing import Optional

import loguru

# from dotenv import load_dotenv
from pydantic import BaseModel

from apropos.src.core.lms.vendors.groq_api import GroqAPIProvider
from apropos.src.core.lms.vendors.openai_api import OpenAIAPIProvider
from typing import Literal

logger = loguru.logger


class JsonDebuggerResponse(BaseModel):
    what_the_error_is: str
    fixed_json_string: str

def get_provider_and_model(provider: Literal["groq", "openai"]):
    if provider == "groq":
        return GroqAPIProvider(force_structured_output=True), "llama3-70b-8192"
    elif provider == "openai":
        return OpenAIAPIProvider(force_structured_output=True), "gpt-4o-mini"
    else:
        raise ValueError("Invalid provider")

def get_messages(
    error_log: list[str],
    string_trying_to_be_parsed_as_json: str,
    response_model: Optional[BaseModel],
):
    response_model_snippet = ""
    if response_model:
        response_model_snippet = f"\nThis json should adhere to the following standard derived from a Pydantic BaseModel: \n\n{response_model.schema()}.\n If you believe the provided response does not contain enough information to be reformatted into the correct JSON, do not make up information, but instead include ESCALATE in your response."
    system_message = f"""
# Instructions
You will be given a JSON string that is not properly formatted, together with the traceback generated by invoking json.loads on the string. Your task is to identify the error in the JSON string and correct it so that it can be parsed successfully.
{response_model_snippet}
Respond ONLY with the corrected JSON string.

# Formatting
Format your response like:
```json
$YOUR_CORRECTED_JSON
```

"""
    user_message = f"""
# Error Log (Last is most recent)
{error_log[-5:]}

# JSON String
{string_trying_to_be_parsed_as_json}

Your correction: """
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_message},
    ]  
    return messages

def parse_outputs(
    fixed_json_string_raw: str,
):
    if "ESCALATE" in fixed_json_string_raw:
        return "ESCALATE"
    if "```json" not in fixed_json_string_raw and "```" in fixed_json_string_raw:
        fixed_json_string_raw = fixed_json_string_raw.split("```")[1].split("```")[
            0
        ]
    elif (
        "```json" not in fixed_json_string_raw
        and "```" not in fixed_json_string_raw
    ):
        fixed_json_string = fixed_json_string_raw
    else:
        fixed_json_string = fixed_json_string_raw.split("```json")[1].split("```")[
            0
        ]
    return fixed_json_string
# TODO: add support for local models?
async def json_debugger_async(
    string_trying_to_be_parsed_as_json: str,
    response_model: Optional[BaseModel],
    provider_name: Literal["groq", "openai"] = "openai",
):
    provider, model = get_provider_and_model(provider_name)
    succeeded = False
    valid_json = None
    patience = 6
    error_log = []

    try:
        json.loads(string_trying_to_be_parsed_as_json)
        succeeded = True
        return json.loads(string_trying_to_be_parsed_as_json)
    except:
        error_log.append("Error in healing: " + traceback.format_exc())
        patience -= 1
    while not succeeded:
        messages = get_messages(error_log, string_trying_to_be_parsed_as_json, response_model)
        fixed_json_string_raw = await provider.async_chat_completion(
            messages=messages,
            model=model,
            temperature=0.0,
            max_tokens=1000,
        )
        fixed_json_string = parse_outputs(fixed_json_string_raw)
        try:
            json.loads(fixed_json_string)
            succeeded = True
            valid_json = fixed_json_string
        except Exception as e:
            error_log.append(
                "Attempted fix: \n\n"
                + fixed_json_string
                + "\n\nError: "
                + traceback.format_exc()
            )
            patience -= 1
        if patience == 0:
            break
    assert succeeded, "Failed to fix JSON string"
    assert not valid_json is None, "Failed to fix JSON string"
    return json.loads(valid_json)


def json_debugger_sync(
    string_trying_to_be_parsed_as_json: str,
    response_model: Optional[BaseModel],
    provider_name: Literal["groq", "openai"] = "openai",
):
    provider, model = get_provider_and_model(provider_name)
    succeeded = False
    valid_json = None
    patience = 6
    error_log = []

    try:
        json.loads(string_trying_to_be_parsed_as_json)
        succeeded = True
        return json.loads(string_trying_to_be_parsed_as_json)
    except:
        error_log.append("Error in healing: " + traceback.format_exc())
        patience -= 1
    while not succeeded:
        messages = get_messages(error_log, string_trying_to_be_parsed_as_json, response_model)
        fixed_json_string_raw = provider.sync_chat_completion(
            messages=messages,
            model=model,
            temperature=0.0,
            max_tokens=1000,
        )
        fixed_json_string = parse_outputs(fixed_json_string_raw)
        try:
            json.loads(fixed_json_string)
            succeeded = True
            valid_json = fixed_json_string
        except Exception as e:
            error_log.append(
                "Attempted fix: \n\n"
                + fixed_json_string
                + "\n\nError: "
                + traceback.format_exc()
            )
            patience -= 1
        if patience == 0:
            break
    assert succeeded, "Failed to fix JSON string"
    assert not valid_json is None, "Failed to fix JSON string"
    return json.loads(valid_json)
