"""
LangGraph-based Car Finder Chatbot

This implementation uses:
- Template-based SQL queries (no SQL generation)
- LangGraph for agent orchestration
- Tool-based architecture for security
- State management for conversation flow
"""

from typing import TypedDict, Annotated, Optional, Literal
from operator import add
import sqlite3
import os
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_core.tools import tool
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field

# Load environment variables
load_dotenv()

# Validate API key
api_key = os.environ.get("OPENAI_API_KEY")
if not api_key:
    raise ValueError("OPENAI_API_KEY environment variable is not set")

# Initialize LLM
llm = ChatOpenAI(model="gpt-4o", temperature=0.7, api_key=api_key)

# Load database schema with error handling
try:
    with open('database_schema.txt', 'r') as f:
        SCHEMA_DESCRIPTION = f.read()
except FileNotFoundError:
    raise FileNotFoundError("database_schema.txt not found. Please ensure it exists in the current directory.")

# Constants
MIN_RESULTS = 1
MAX_RESULTS = 20
DB_PATH = 'cars.db'


# ============================================================================
# PYDANTIC MODELS FOR STRUCTURED TOOL INPUTS
# ============================================================================

class SearchParameters(BaseModel):
    """Parameters for searching cars using template-based SQL"""
    min_price: Optional[int] = Field(None, description="Minimum price in USD", ge=0, le=100000)
    max_price: Optional[int] = Field(None, description="Maximum price in USD", ge=0, le=100000)
    fuel_type: Optional[Literal["Gasoline", "Diesel", "Electric", "Hybrid", "Plug-in Hybrid"]] = Field(None, description="Type of fuel")
    is_suv: Optional[bool] = Field(None, description="True for SUVs, False for sedans/coupes")
    min_seating: Optional[int] = Field(None, description="Minimum seating capacity", ge=4, le=8)
    max_seating: Optional[int] = Field(None, description="Maximum seating capacity", ge=4, le=8)
    drivetrain: Optional[Literal["FWD", "RWD", "AWD", "4WD"]] = Field(None, description="Drive system")
    min_fuel_efficiency_city: Optional[float] = Field(None, description="Minimum city MPG", ge=0)
    min_cargo_space: Optional[int] = Field(None, description="Minimum cargo space in cubic feet", ge=0)
    has_sunroof: Optional[bool] = Field(None, description="Must have sunroof")
    has_leather_seats: Optional[bool] = Field(None, description="Must have leather seats")
    has_navigation: Optional[bool] = Field(None, description="Must have navigation system")
    has_backup_camera: Optional[bool] = Field(None, description="Must have backup camera")
    min_safety_rating: Optional[float] = Field(None, description="Minimum safety rating", ge=0, le=5)


# ============================================================================
# SQL TEMPLATE-BASED TOOLS (Secure, No SQL Generation)
# ============================================================================

@tool
def search_cars(
    min_price: Optional[int] = None,
    max_price: Optional[int] = None,
    fuel_type: Optional[str] = None,
    is_suv: Optional[bool] = None,
    min_seating: Optional[int] = None,
    max_seating: Optional[int] = None,
    drivetrain: Optional[str] = None,
    min_fuel_efficiency_city: Optional[float] = None,
    min_cargo_space: Optional[int] = None,
    has_sunroof: Optional[bool] = None,
    has_leather_seats: Optional[bool] = None,
    has_navigation: Optional[bool] = None,
    has_backup_camera: Optional[bool] = None,
    min_safety_rating: Optional[float] = None,
) -> dict:
    """
    Search for cars using a secure template-based SQL query.

    Returns a dictionary with:
    - count: number of cars found
    - cars: list of matching cars (limited to first 20)
    - status: 'too_few', 'good', or 'too_many'
    """
    # Build WHERE clause from parameters using template
    conditions = []
    params = []

    if min_price is not None:
        conditions.append("price >= ?")
        params.append(min_price)

    if max_price is not None:
        conditions.append("price <= ?")
        params.append(max_price)

    if fuel_type is not None:
        conditions.append("fuel_type = ?")
        params.append(fuel_type)

    if is_suv is not None:
        conditions.append("is_suv = ?")
        params.append(1 if is_suv else 0)

    if min_seating is not None:
        conditions.append("seating_capacity >= ?")
        params.append(min_seating)

    if max_seating is not None:
        conditions.append("seating_capacity <= ?")
        params.append(max_seating)

    if drivetrain is not None:
        conditions.append("drivetrain = ?")
        params.append(drivetrain)

    if min_fuel_efficiency_city is not None:
        conditions.append("fuel_efficiency_city >= ?")
        params.append(min_fuel_efficiency_city)

    if min_cargo_space is not None:
        conditions.append("cargo_space >= ?")
        params.append(min_cargo_space)

    if has_sunroof is not None:
        conditions.append("has_sunroof = ?")
        params.append(1 if has_sunroof else 0)

    if has_leather_seats is not None:
        conditions.append("has_leather_seats = ?")
        params.append(1 if has_leather_seats else 0)

    if has_navigation is not None:
        conditions.append("has_navigation = ?")
        params.append(1 if has_navigation else 0)

    if has_backup_camera is not None:
        conditions.append("has_backup_camera = ?")
        params.append(1 if has_backup_camera else 0)

    if min_safety_rating is not None:
        conditions.append("safety_rating >= ?")
        params.append(min_safety_rating)

    # Build SQL query using template
    where_clause = " AND ".join(conditions) if conditions else "1=1"
    query = f"SELECT * FROM cars WHERE {where_clause} ORDER BY price LIMIT 21"  # Get 21 to check if > 20

    try:
        # Use context manager to ensure connection is closed
        with sqlite3.connect(DB_PATH) as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute(query, params)
            results = cursor.fetchall()

            # Convert to list of dicts
            cars = [dict(row) for row in results]
            count = len(cars)

            # Determine status
            if count < MIN_RESULTS:
                status = "too_few"
            elif count > MAX_RESULTS:
                status = "too_many"
                cars = cars[:MAX_RESULTS]  # Limit results
            else:
                status = "good"

            return {
                "count": count,
                "cars": cars,
                "status": status,
                "params_used": {k: v for k, v in [
                    ("min_price", min_price),
                    ("max_price", max_price),
                    ("fuel_type", fuel_type),
                    ("is_suv", is_suv),
                    ("min_seating", min_seating),
                    ("max_seating", max_seating),
                    ("drivetrain", drivetrain),
                    ("min_fuel_efficiency_city", min_fuel_efficiency_city),
                    ("min_cargo_space", min_cargo_space),
                    ("has_sunroof", has_sunroof),
                    ("has_leather_seats", has_leather_seats),
                    ("has_navigation", has_navigation),
                    ("has_backup_camera", has_backup_camera),
                    ("min_safety_rating", min_safety_rating),
                ] if v is not None}
            }

    except sqlite3.Error as e:
        return {
            "count": 0,
            "cars": [],
            "status": "error",
            "error": str(e)
        }


@tool
def count_cars_only(
    min_price: Optional[int] = None,
    max_price: Optional[int] = None,
    fuel_type: Optional[str] = None,
    is_suv: Optional[bool] = None,
    min_seating: Optional[int] = None,
    max_seating: Optional[int] = None,
    drivetrain: Optional[str] = None,
    min_fuel_efficiency_city: Optional[float] = None,
    min_cargo_space: Optional[int] = None,
    has_sunroof: Optional[bool] = None,
    has_leather_seats: Optional[bool] = None,
    has_navigation: Optional[bool] = None,
    has_backup_camera: Optional[bool] = None,
    min_safety_rating: Optional[float] = None,
) -> dict:
    """
    Count how many cars match the criteria without returning full results.
    Useful for checking if we need to refine search parameters.
    """
    # Build WHERE clause from parameters
    conditions = []
    params = []

    if min_price is not None:
        conditions.append("price >= ?")
        params.append(min_price)

    if max_price is not None:
        conditions.append("price <= ?")
        params.append(max_price)

    if fuel_type is not None:
        conditions.append("fuel_type = ?")
        params.append(fuel_type)

    if is_suv is not None:
        conditions.append("is_suv = ?")
        params.append(1 if is_suv else 0)

    if min_seating is not None:
        conditions.append("seating_capacity >= ?")
        params.append(min_seating)

    if max_seating is not None:
        conditions.append("seating_capacity <= ?")
        params.append(max_seating)

    if drivetrain is not None:
        conditions.append("drivetrain = ?")
        params.append(drivetrain)

    if min_fuel_efficiency_city is not None:
        conditions.append("fuel_efficiency_city >= ?")
        params.append(min_fuel_efficiency_city)

    if min_cargo_space is not None:
        conditions.append("cargo_space >= ?")
        params.append(min_cargo_space)

    if has_sunroof is not None:
        conditions.append("has_sunroof = ?")
        params.append(1 if has_sunroof else 0)

    if has_leather_seats is not None:
        conditions.append("has_leather_seats = ?")
        params.append(1 if has_leather_seats else 0)

    if has_navigation is not None:
        conditions.append("has_navigation = ?")
        params.append(1 if has_navigation else 0)

    if has_backup_camera is not None:
        conditions.append("has_backup_camera = ?")
        params.append(1 if has_backup_camera else 0)

    if min_safety_rating is not None:
        conditions.append("safety_rating >= ?")
        params.append(min_safety_rating)

    where_clause = " AND ".join(conditions) if conditions else "1=1"
    query = f"SELECT COUNT(*) as count FROM cars WHERE {where_clause}"

    try:
        with sqlite3.connect(DB_PATH) as conn:
            cursor = conn.cursor()
            cursor.execute(query, params)
            count = cursor.fetchone()[0]

            return {
                "count": count,
                "status": "too_few" if count < MIN_RESULTS else "too_many" if count > MAX_RESULTS else "good"
            }

    except sqlite3.Error as e:
        return {
            "count": 0,
            "status": "error",
            "error": str(e)
        }


# ============================================================================
# LANGGRAPH STATE DEFINITION
# ============================================================================

class ConversationState(TypedDict):
    """State for the conversation graph"""
    messages: Annotated[list, add]  # Conversation history
    search_params: Optional[dict]  # Current search parameters
    search_results: Optional[dict]  # Results from last search
    iteration_count: int  # Number of search iterations
    user_satisfied: bool  # Whether user is satisfied
    requires_search: bool  # Whether we need to search


# ============================================================================
# LANGGRAPH NODES
# ============================================================================

def gather_requirements(state: ConversationState) -> ConversationState:
    """
    Node: Gather requirements from user and determine search parameters.
    This node uses LLM to extract search criteria from conversation.
    """
    messages = state["messages"]

    system_prompt = f"""You are a friendly car shopping advisor. Analyze the conversation and extract search parameters.

{SCHEMA_DESCRIPTION}

Based on the user's requirements, determine:
1. What search parameters should be used (price range, fuel type, SUV/sedan, features, etc.)
2. Whether we have enough information to search (set requires_search=true)
3. Whether the user is satisfied with current results (set user_satisfied=true)

If the user just greeted you or asked a general question, be friendly and ask what they're looking for.
If previous search returned too few/many results, suggest adjustments.

You have access to these tools:
- search_cars: Search for cars with specific parameters
- count_cars_only: Check count before full search

IMPORTANT: Only use tools when you have concrete search criteria.
For greetings or clarifying questions, just respond conversationally without calling tools."""

    # Get LLM response with tools
    llm_with_tools = llm.bind_tools([search_cars, count_cars_only])
    response = llm_with_tools.invoke([SystemMessage(content=system_prompt)] + messages)

    # Update state with response
    new_state = {
        "messages": [response],
        "requires_search": bool(response.tool_calls),
        "iteration_count": state.get("iteration_count", 0)
    }

    # Check if user expressed satisfaction
    if "satisfied" in response.content.lower() or "perfect" in response.content.lower():
        new_state["user_satisfied"] = True

    return new_state


def execute_search(state: ConversationState) -> ConversationState:
    """
    Node: Execute the search using tool calls from the LLM.
    """
    messages = state["messages"]
    last_message = messages[-1]

    # Execute tool calls
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        tool_node = ToolNode([search_cars, count_cars_only])
        result = tool_node.invoke(state)

        # Extract search results
        tool_messages = result["messages"]
        if tool_messages:
            # Parse the tool response
            tool_response = tool_messages[-1]
            if hasattr(tool_response, 'content'):
                import json
                try:
                    search_results = json.loads(tool_response.content)
                except:
                    search_results = {"error": "Failed to parse tool response"}
            else:
                search_results = {}
        else:
            search_results = {}

        return {
            "messages": tool_messages,
            "search_results": search_results,
            "iteration_count": state.get("iteration_count", 0) + 1
        }

    return {"iteration_count": state.get("iteration_count", 0)}


def present_results(state: ConversationState) -> ConversationState:
    """
    Node: Present search results to user and provide guidance.
    """
    search_results = state.get("search_results", {})
    messages = state["messages"]

    system_prompt = """You are presenting car search results to the user.

Based on the search results:
- If status is 'good' (1-20 cars): Present the results enthusiastically and ask if they want details
- If status is 'too_few' (<1 cars): Suggest broadening criteria (increase price range, consider more fuel types, etc.)
- If status is 'too_many' (>20 cars): Suggest narrowing criteria (set budget, choose specific features, etc.)
- If status is 'error': Apologize and ask them to rephrase

Be conversational, helpful, and guide the user toward finding their perfect car."""

    # Create context message with results
    context = f"\nSearch Results: {search_results}"

    response = llm.invoke([
        SystemMessage(content=system_prompt),
        *messages,
        HumanMessage(content=context)
    ])

    return {"messages": [response]}


def should_continue(state: ConversationState) -> Literal["execute_search", "present_results", "end"]:
    """
    Conditional edge: Determine next step in the graph.
    """
    # Check if user is satisfied
    if state.get("user_satisfied", False):
        return "end"

    # Check if we need to execute search
    if state.get("requires_search", False) and not state.get("search_results"):
        return "execute_search"

    # If we have search results, present them
    if state.get("search_results"):
        return "present_results"

    # Otherwise, end (probably a greeting or farewell)
    return "end"


# ============================================================================
# BUILD LANGGRAPH
# ============================================================================

def build_graph() -> StateGraph:
    """Build the LangGraph workflow"""

    workflow = StateGraph(ConversationState)

    # Add nodes
    workflow.add_node("gather_requirements", gather_requirements)
    workflow.add_node("execute_search", execute_search)
    workflow.add_node("present_results", present_results)

    # Set entry point
    workflow.set_entry_point("gather_requirements")

    # Add conditional edges
    workflow.add_conditional_edges(
        "gather_requirements",
        should_continue,
        {
            "execute_search": "execute_search",
            "present_results": "present_results",
            "end": END
        }
    )

    # After executing search, present results
    workflow.add_edge("execute_search", "present_results")

    # After presenting results, go back to gather requirements for next turn
    workflow.add_edge("present_results", END)

    return workflow.compile()


# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

def format_car_display(car: dict) -> str:
    """Format a single car for display"""
    return f"""
{car['brand']} {car['model_name']} ({car['year']})
  Price: ${car['price']:,}
  Type: {'SUV' if car['is_suv'] else 'Sedan/Coupe'}
  Fuel: {car['fuel_type']}
  Seats: {car['seating_capacity']} | Cargo: {car['cargo_space']} cu ft
  Drivetrain: {car['drivetrain']} | Transmission: {car['transmission']}
  Features: {'Sunroof, ' if car['has_sunroof'] else ''}{'Leather, ' if car['has_leather_seats'] else ''}{'Navigation, ' if car['has_navigation'] else ''}{'Backup Camera' if car['has_backup_camera'] else ''}
"""


# ============================================================================
# MAIN CHATBOT LOOP
# ============================================================================

def main():
    """Main chatbot loop using LangGraph"""
    print("=" * 60)
    print("Welcome to Car Finder Chatbot (LangGraph Edition)")
    print("=" * 60)
    print("Tell me what kind of car you're looking for!")
    print("Type 'quit' to exit.\n")

    # Build the graph
    app = build_graph()

    # Initialize conversation state
    conversation_state = {
        "messages": [],
        "search_params": None,
        "search_results": None,
        "iteration_count": 0,
        "user_satisfied": False,
        "requires_search": False
    }

    while True:
        # Get user input
        user_input = input("You: ").strip()

        if user_input.lower() in ['quit', 'exit', 'bye']:
            print("\nThanks for using Car Finder! Goodbye! 👋")
            break

        if not user_input:
            continue

        # Add user message to state
        conversation_state["messages"].append(HumanMessage(content=user_input))

        # Run the graph
        try:
            result = app.invoke(conversation_state)

            # Update conversation state
            conversation_state = result

            # Display assistant's response
            if result["messages"]:
                last_message = result["messages"][-1]
                if hasattr(last_message, 'content'):
                    print(f"\nAssistant: {last_message.content}\n")

                    # If we have good search results, display them
                    if result.get("search_results", {}).get("status") == "good":
                        cars = result["search_results"].get("cars", [])
                        if cars:
                            print("=" * 60)
                            print("✅ MATCHING CARS:")
                            print("=" * 60)
                            for car in cars:
                                print(format_car_display(car))
                            print("=" * 60)
                            print()

            # Check if conversation should end
            if result.get("user_satisfied", False):
                print("\nThank you for using Car Finder! Have a great day! 👋")
                break

        except Exception as e:
            print(f"\n❌ Error: {str(e)}")
            print("Let's try again. Please rephrase your request.\n")
            # Remove the failed message
            conversation_state["messages"] = conversation_state["messages"][:-1]


if __name__ == "__main__":
    main()