import os
import json
from datetime import datetime, timedelta
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_core.tools import Tool
from langchain.agents import create_openai_tools_agent, AgentExecutor
from langchain.prompts import (
    PromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
    ChatPromptTemplate,
)
import pytz
import uvicorn

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise ValueError("OPENAI_API_KEY is not set.")

app = FastAPI()

# Database connection
db_instance = SQLDatabase.from_uri(
    "mssql+pyodbc://PMOlogbook_Usr:aPMO86#iaxh@103.239.89.100:21433/PMO360_DB?driver=ODBC Driver 17 for SQL Server",
    include_tables=["Logbook", "LogbookHistory", "User", "tbl_Sprint", "tblMs_Projects"]
)

# Helper: get current week range
def get_week_range():
    today = datetime.now(pytz.timezone("Asia/Kolkata"))
    start = today - timedelta(days=today.weekday())
    end = start + timedelta(days=6)
    return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")

# Helper: SQL Query Builder
def build_sql(first_name: str, last_name: str = None, start_date: str = None, end_date: str = None) -> str:
    if not start_date or not end_date:
        start_date, end_date = get_week_range()

    name_filter = f"u.FirstName = '{first_name}'"
    if last_name:
        name_filter += f" AND u.LastName = '{last_name}'"

    return f"""
    SELECT 
        COALESCE(u.FirstName + ' ' + u.LastName, 'Unknown') AS Name,
        lh.LogbookId,
        lh.CreatedDate,
        p.ProjectName,
        s.Track,
        s.SubTrack,
        lh.Description,
        lh.Comment
    FROM LogbookHistory lh
    LEFT JOIN [User] u ON lh.AssisgnTo = u.UserId
    LEFT JOIN [Logbook] l ON lh.LogbookId = l.LogbookId
    LEFT JOIN [tblMs_Projects] p ON l.ProjectName = p.ProjectID
    LEFT JOIN [tbl_Sprint] s ON lh.LogbookId = s.LogbookId
    WHERE {name_filter}
      AND lh.CreatedDate BETWEEN '{start_date}' AND '{end_date}'
    ORDER BY lh.CreatedDate ASC
    """

# SQL tool for agent
query_weekly_progress = Tool(
    name="query_weekly_progress",
    func=lambda query: db_instance.run(query) if query.lower().strip().startswith("select") else "Only SELECT queries are allowed.",
    description="Runs a SQL SELECT query to retrieve user progress report."
)

# System prompt for clean HTML output
system_prompt = SystemMessagePromptTemplate(
    prompt=PromptTemplate(
        input_variables=[],
        template="""
You are a weekly reporting agent for PMO360.
Return a clean HTML <table> using SQL output:
- Fields: Name, LogbookId, Date, Project Name, Track, Sub Track, Description, Comment
- Include headers
- Align text to the left
- Add table borders using a constant CSS class
- No extra explanation, only output the HTML table
"""
    )
)

# Human prompt
human_prompt = HumanMessagePromptTemplate(
    prompt=PromptTemplate(input_variables=["input"], template="{input}")
)

# Chat prompt
chat_prompt = ChatPromptTemplate.from_messages([
    system_prompt,
    MessagesPlaceholder(variable_name="chat_history", optional=True),
    human_prompt,
    MessagesPlaceholder(variable_name="agent_scratchpad")
])

# Setup LLM + agent executor
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key=OPENAI_API_KEY)
agent = create_openai_tools_agent(llm=llm, tools=[query_weekly_progress], prompt=chat_prompt)
agent_executor = AgentExecutor(agent=agent, tools=[query_weekly_progress], verbose=True)

# Name extractor LLM
name_extractor = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key=OPENAI_API_KEY)
name_prompt = PromptTemplate(
    input_variables=["text"],
    template="""
Extract the first name, last name, and optional date range from the following text. Respond ONLY with a dictionary (no strings or quotes):
Examples:
"Show Rahul Kumar's report from April 10 to April 17" → {{"first_name": "Rahul", "last_name": "Kumar", "start_date": "2025-04-10", "end_date": "2025-04-17"}}
"DPR for Anirban this week" → {{"first_name": "Anirban"}}
"DPR is daily progress report"
Prompt: {text}
"""
)

# FastAPI model for prompt input
class PromptRequest(BaseModel):
    prompt: str

# API endpoint
@app.post("/weekly-dpr")
async def get_prompt_based_dpr(request: PromptRequest):
    try:
        prompt_text = request.prompt.strip()
        name_response = name_extractor.invoke(name_prompt.format_prompt(text=prompt_text).to_string())
        print("LLM raw output:", name_response.content)

        # Parse extracted names
        parsed = name_response.content.strip()
        if parsed.startswith("\"") or parsed.startswith("'"):
            parsed = parsed[1:-1]
        name_json = json.loads(parsed.replace("'", '"'))

        first = name_json.get("first_name")
        last = name_json.get("last_name")
        start_date = name_json.get("start_date")
        end_date = name_json.get("end_date")

        if not first:
            raise ValueError("No valid first name extracted.")

        sql = build_sql(first, last, start_date, end_date)
        context = f"Generate HTML weekly report for: {first} {last or ''} using this query:\n{sql}. Table must include Name, LogbookId, CreatedDate, ProjectName, Track, SubTrack, Description, Comment, and use a CSS class for bordered tables."

        response = agent_executor.invoke({"input": context})
        return response.get("output", "No output generated.")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"DPR Error: {e}")

# Run the app
if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)