from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
import pandas as pd
from sqlalchemy import create_engine
from datetime import datetime, timedelta
from langchain_core.prompts import PromptTemplate 
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace

import os
import gradio as gr
#from openai import OpenAI
from huggingface_hub import InferenceClient
from app.hybrid_rag import HybridJiraRAG
from langchain_core.output_parsers import StrOutputParser

# Sample Jira data structure
jira_data = {
    'ticket_id': ['PROJ-1', 'PROJ-2', 'PROJ-3', 'PROJ-4'],
    'summary': ['Bug in login', 'Feature request', 'Performance issue', 'Security bug'],
    'status': ['Closed', 'Open', 'In Progress', 'Closed'],
    'priority': ['P1', 'P3', 'P2', 'P1'],
    'severity': ['Critical', 'Low', 'Medium', 'Critical'],
    'created_date': ['2024-01-01', '2024-01-05', '2024-01-10', '2024-01-15'],
    'closed_date': ['2024-01-03', None, None, '2024-01-16'],
    'resolution_time_hours': [48, None, None, 24],
    'assignee': ['john', 'jane', 'bob', 'alice']
}

df = pd.DataFrame(jira_data)
df['created_date'] = pd.to_datetime(df['created_date'])
df['closed_date'] = pd.to_datetime(df['closed_date'])

# Save to SQLite
engine = create_engine("sqlite:///jira.db")
df.to_sql("tickets", engine, if_exists="replace", index=False)

# Create agent with table description
db = SQLDatabase(engine)

print("Tables:", db.get_table_names())
print("\nSchema:", db.get_table_info())
print("\nSample Data:", db.run("SELECT * FROM tickets LIMIT 5"))

# Add table descriptions for better context
table_info = """
The 'tickets' table contains Jira ticket data with columns:
- ticket_id: Unique ticket identifier (e.g., PROJ-123)
- summary: Brief description of the ticket
- status: Current status (Open, In Progress, Closed)
- priority: Priority level (P1=Highest, P2=High, P3=Medium, P4=Low)
- severity: Severity level (Critical, High, Medium, Low)
- created_date: When ticket was created
- closed_date: When ticket was closed (NULL if still open)
- resolution_time_hours: Time taken to resolve in hours
- assignee: Person assigned to ticket
"""

#llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

#llm = InferenceClient(api_key=settings.HF_TOKEN)
endpoint = HuggingFaceEndpoint(
    repo_id="meta-llama/Llama-3.1-8B-Instruct",
    huggingfacehub_api_token=os.environ["HF_TOKEN"],
    temperature=0.1,
    max_new_tokens=512
)

llm = ChatHuggingFace(llm=endpoint)

agent = create_sql_agent(
    llm, 
    db=db, 
    verbose=True,
    #agent_type="openai-tools"
    agent_type="zero-shot-react-description",
    handle_parsing_errors=True
)

# Query examples
questions = [
    "What is the average resolution time?",
    "How many tickets are open vs closed?",
    "Show distribution of tickets by severity",
    "Which assignee has the most P1 tickets?",
    "How many critical tickets were resolved in less than 48 hours?",
]

for q in questions:
    print(f"\n{'='*60}")
    print(f"Q: {q}")
    print(f"{'='*60}")
    result = agent.invoke(q)
    print(f"A: {result['output']}\n")



# After getting SQL results, format them nicely
def ask_with_formatting(question: str):
    # Generate and execute SQL
    sql = sql_chain.invoke({"question": question})
    raw_result = execute_query.invoke(sql)
    
    # Format result in natural language
    format_prompt = PromptTemplate(
        template="""Given the question and SQL result, provide a clear natural language answer.
        
        Question: {question}
        SQL Result: {result}

        Natural language answer:""",
                input_variables=["question", "result"]
    )
    
    #format_chain = LLMChain(llm=llm, prompt=format_prompt)
    format_chain = prompt | llm | StrOutputParser()
    formatted = format_chain.invoke({
        "question": question,
        "result": raw_result
    })
    
    return formatted['text']

# Usage
print(ask_with_formatting("What's the average resolution time for P1 tickets?"))
# Output: "The average resolution time for P1 priority tickets is 36 hours, or approximately 1.5 days."


# Build Gradio UI
demo = gr.Interface(
    fn=ask_llm,
    inputs=gr.Textbox(lines=3, label="Ask the AI"),
    outputs=gr.Textbox(label="Response"),
    title="HF Inference Client LLM Demo",
    description="Powered by HuggingFace InferenceClient SDK."
)

demo.launch()