import os
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
import json  
import openai
from langsmith.wrappers import wrap_openai
from langsmith import traceable
from langchain import hub
import re

load_dotenv()

# --------------------------------------------------
# Utility Functions
# --------------------------------------------------

client = wrap_openai(openai.OpenAI())

chatModel_json = ChatOpenAI(
    model="gpt-4o",
    model_kwargs={"response_format": {"type": "json_object"}},
)

chatModel_no_json = ChatOpenAI(
    model="gpt-4o"
)

async def parse_json_with_retries(json_string, max_attempts=3):
    for attempt in range(max_attempts):
        try:
            return json.loads(json_string)
        except json.JSONDecodeError as error:
            if attempt == max_attempts - 1:
                raise error

@traceable
async def invoke_agent(runnable, params):
    content = runnable.invoke(params)
    if hasattr(content, 'content'):
        content = content.content
    return content

# --------------------------------------------------
# Translation Function
# --------------------------------------------------

@traceable
async def translate_text(text, target_language):
    
    try:
        first_prompt = hub.pull("getgloby/agent-1-get-info")
        print("First prompt pulled successfully.")
    except Exception as e:
        print(f"Error al obtener el primer agente: {e}")
        return

    first_runnable = first_prompt | chatModel_no_json

    try:
        first_agent_output = await invoke_agent(first_runnable, {
            "originalText": text,
            "targetLanguage": target_language,
        })
        print(f"First agent output: {first_agent_output}")
    except Exception as e:
        print(f"Error invoking first agent: {e}")
        print(f"Error details: {e}")
        return

    translations = []

    for i in range(3):
        try:
            second_prompt = hub.pull("getgloby/agent-2-translate")
            print(f"Second prompt pulled successfully for iteration {i + 1}.")
        except Exception as e:
            print(f"Error al obtener el segundo agente en la iteración {i + 1}: {e}")
            continue

        second_runnable = second_prompt | chatModel_json

        try:
            second_agent_output = await invoke_agent(second_runnable, {
                "originalText": text,
                "firstAgentOutput": first_agent_output,
                "targetLanguage": target_language,
            })
            print(f"Second agent output (iteration {i + 1}): {second_agent_output}")
        except Exception as e:
            print(f"Error invoking second agent in iteration {i + 1}: {e}")
            print(f"Error details: {e}")
            continue

        try:
            parsed_second_agent_output = await parse_json_with_retries(second_agent_output)
            translation = parsed_second_agent_output["translation"]
            print(f"Parsed second agent output (iteration {i + 1}): {translation}")
        except Exception as e:
            print(f"Error parsing second agent output in iteration {i + 1}: {e}")
            print(f"Error details: {e}")
            continue

        try:
            third_prompt = hub.pull("getgloby/agent-3-improve-translation")
            print(f"Third prompt pulled successfully for iteration {i + 1}.")
        except Exception as e:
            print(f"Error al obtener el tercer agente en la iteración {i + 1}: {e}")
            continue

        third_runnable = third_prompt | chatModel_json

        try:
            third_agent_output = await invoke_agent(third_runnable, {
                "originalText": text,
                "firstAgentOutput": first_agent_output,
                "secondAgentOutput": translation,
            })
            print(f"Third agent output (iteration {i + 1}): {third_agent_output}")
        except Exception as e:
            print(f"Error invoking third agent in iteration {i + 1}: {e}")
            print(f"Error details: {e}")
            continue

        try:
            third_agent_output_json = await parse_json_with_retries(third_agent_output)
            optimized_translation = third_agent_output_json["optimizedTranslation"]
            print(f"Optimized translation (iteration {i + 1}): {optimized_translation}")
        except Exception as e:
            print(f"Error parsing third agent output in iteration {i + 1}: {e}")
            print(f"Error details: {e}")
            continue

        translations.append(optimized_translation)

    try:
        final_prompt = hub.pull("getgloby/agent-4-final-translation")
        print("Final prompt pulled successfully.")
    except Exception as e:
        print(f"Error al obtener el agente final: {e}")
        return

    final_runnable = final_prompt | chatModel_json

    try:
        final_agent_output = await invoke_agent(final_runnable, {
            "originalText": text,
            "firstAgentOutput": first_agent_output,  # Cambiado de contextInfo a firstAgentOutput
            "translation1": translations[0],
            "translation2": translations[1],
            "translation3": translations[2],
            "targetLanguage": target_language,
        })
        print(f"Final agent output: {final_agent_output}")
    except Exception as e:
        print(f"Error invoking final agent: {e}")
        print(f"Error details: {e}")
        return

    try:
        final_agent_output_json = await parse_json_with_retries(final_agent_output)
        final_translation = final_agent_output_json["finalTranslation"]
        print(f"Final translation: {final_translation}")
    except Exception as e:
        print(f"Error parsing final agent output: {e}")
        print(f"Error details: {e}")
        return

    return final_translation