import gradio as gr
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.agents import AgentExecutor, create_structured_chat_agent
from langchain_community.callbacks import get_openai_callback
import os
import time
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain import hub

# Load environment variables
API_KEY = os.getenv("OPENAI_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
PASSWORD = os.getenv("APP_PASSWORD")

if not API_KEY:
    raise ValueError("Please set the OPENAI_API_KEY environment variable.")

if not TAVILY_API_KEY:
    raise ValueError("Please set the TAVILY_API_KEY environment variable.")

if not PASSWORD:
    raise ValueError("Please set the APP_PASSWORD environment variable.")

# Define the enhanced prompt template with instructions to include sources and Markdown formatting
prompt_template = PromptTemplate(
    input_variables=["marka_model", "automobil", "question"],
    template="""
Respond in systematically cleaned-up Markdown format and only in Russian: You are a virtual technical expert with years of experience in vehicle repair and maintenance. Your responses are intended for professional mechanics and contain only reliable and safe instructions that adhere to best practices and manufacturer recommendations.

We are currently discussing repair work for the {marka_model} model vehicle {automobil}. Please provide your expert opinion on the following question: {question}

### When crafting your response, ensure the following:
1. **Step-by-Step Repair Instructions**:
   - Provide a clear, sequential plan of action to resolve the issue.
   - List necessary tools, parts, and materials, specifying time and labor estimates.

2. **Technical Specifications and Parameters**:
   - Highlight key technical specifications and manufacturer recommendations, citing sources when possible.

3. **Alternative Solutions**:
   - Include any possible alternative methods that may resolve the issue in a more cost-effective manner, if applicable.

4. **Visual Aids**:
   - Recommend multimedia resources, such as videos, to support the text instructions where applicable.

### Response Formatting:
- Begin with a table outlining each repair step. In one column, include the name of the operation, and in the other, the time in labor hours. At the bottom of the table, display the total repair time in labor hours.
- Follow the table with a detailed, step-by-step description of the repair operations.

Balance confidence with caution to ensure the safety and reliability of the completed repair.
"""
)

# Initialize components with updated classes
llm = ChatOpenAI(
    openai_api_key=API_KEY,
    model_name="gpt-4o-mini",
    temperature=0,
)

tools = [TavilySearchResults(max_results=5)]
prompt = prompt_template

# Initialize the structured chat agent
prompt = hub.pull("hwchase17/structured-chat-agent")
agent = create_structured_chat_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools)

def generate_response(marka_model, automobil, question, password):
    """
    Generates a response based on the user's automotive repair question.

    Args:
        marka_model (str): Car make and model.
        automobil (str): Vehicle details.
        question (str): Automotive repair question.
        password (str): Password for accessing the app.

    Returns:
        tuple: Response in Markdown format and performance metrics in JSON.
    """
    if password != PASSWORD:
        return {"error": "Incorrect password. Access denied."}
    
    try:
        with get_openai_callback() as cb:
            start_time = time.perf_counter()
            user_input = {  
                "marka_model": marka_model,
                "automobil": automobil,
                "question": question
            }
            formatted_prompt = prompt_template.format(**user_input)
            response = agent_executor.invoke({"input": formatted_prompt})
            response = llm.invoke(f'Clean up the response and translate it to Russian if needed: {response}').content
            end_time = time.perf_counter()
    
        total_time = end_time - start_time
    
        # Updated cost rates
        costs = {'input': 0.150 / 1e6, 'output': 0.600 / 1e6}
    
        total_cost = cb.prompt_tokens * costs['input'] + cb.completion_tokens * costs['output']
    
        json_data = {
            "total_cost": round(total_cost, 6),
            "total_time_seconds": round(total_time, 4)
        }
    
        return response, json_data
    except Exception as e:
        return {"error": str(e)}, {}

# Define example inputs
examples = [
    [
        "HYUNDAI SOLARIS",
        "HYUNDAI SOLARIS 01.01.2012 0:00:00 бензин 1,4 - (G4FA 4/16) АКПП (A4CF1) FWD - ABS",
        "что нужно снимать при замене заднего сайлентблока переднего нижнего рычага?"
    ],
    [
        "TOYOTA CAMRY",
        "TOYOTA CAMRY 2015 бензин 2.5 - (2AR-FE) АКПП (U760E) FWD",
        "как заменить тормозные колодки на передних колесах?"
    ],
    [
        "FORD FOCUS",
        "FORD FOCUS 2018 дизель 1.5 - (DV5RC) МКПП (B6) FWD",
        "какие инструменты нужны для замены масла в двигателе?"
    ],
    [
        "BMW X5",
        "BMW X5 2020 бензин 3.0 - (B58) АКПП (ZF 8HP) AWD",
        "как проверить уровень жидкости в коробке передач?"
    ],
    [
        "HONDA CIVIC",
        "HONDA CIVIC 2010 бензин 1.8 - (R18A) МКПП (LSD) FWD",
        "как заменить воздушный фильтр двигателя?"
    ]
]

iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Car Make and Model", placeholder="Enter the car make and model"),
        gr.Textbox(label="Vehicle", placeholder="Enter the vehicle"),
        gr.Textbox(label="Your Question", placeholder="Enter your automotive repair question"),
        gr.Textbox(label="Password", placeholder="Enter the password to access the app"),
    ],
    outputs=[
        gr.Markdown(label="Answer"),
        gr.JSON(label="Performance Metrics")
    ],
    title="Automotive Technical Assistant",
    description="Get detailed automotive repair and maintenance instructions",
    examples=examples,
    cache_examples=False,
)

if __name__ == "__main__":
    iface.launch(share=True)