# main.py
from dotenv import load_dotenv
import re
import os
import json
from typing import List

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from models import (
    TravelDataInput,
    PlanResponse,
    ExtractionInput,
    PackageData
)
from image_searcher import ImageSearcher, create_image_queries, extract_destination
from gemini_utils import create_travel_plan_prompt, generate_with_gemini
from rag_utils import get_relevant_plan

load_dotenv()
PEXELS_API_KEY = os.getenv("PEXELS_API_KEY")
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# App setup
app = FastAPI(title="Travel Recommendation API")
image_searcher = ImageSearcher(PEXELS_API_KEY)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/", summary="Root endpoint for health check")
def read_root():
    return {"message": "Travel Recommendation API is running!"}


@app.post("/generate_plan", response_model=PlanResponse, summary="Generate a travel plan")
async def generate_plan_endpoint(travel_input: TravelDataInput):
    travel_data = travel_input.model_dump()
    source = "generated"
    plan_text = None

    # Step 1: RAG
    try:
        relevant_plan = get_relevant_plan(travel_data)
        if relevant_plan:
            plan_text = relevant_plan
            source = "index"
    except Exception as e:
        print(f"RAG check failed: {e}")

    # Step 2: Generate if RAG fails
    if not plan_text:
        try:
            prompt = create_travel_plan_prompt(travel_data)
            plan_text = generate_with_gemini(prompt)
        except Exception as e:
            print(f"Gemini generation failed: {e}")
            raise HTTPException(status_code=500, detail=f"Failed to generate plan: {e}")

    if not plan_text:
        raise HTTPException(status_code=500, detail="Could not retrieve or generate a travel plan.")

    # Step 3: Get destination & images
    try:
        destination = extract_destination(plan_text, travel_data) or "Suggested Destination"
        image_queries = create_image_queries(travel_data, destination)
        images = image_searcher.search_all_sources(image_queries, num_images=6)
    except Exception as e:
        print(f"Image search failed: {e}")
        images = []
        image_queries = []

    return PlanResponse(
        plan_text=plan_text,
        destination_used=destination,
        images=images,
        image_queries_used=image_queries,
        source=source,
    )


@app.post("/extract_package", response_model=PackageData, summary="Extract structured package data from plan text")
async def extract_package_endpoint(data: ExtractionInput):
    extraction_prompt = f"""
Extract ONLY the following fields from the travel plan below and return STRICTLY in raw JSON format. 
No extra explanation, no markdown, no text.

Required JSON format:
{{
  "packageName": string,
  "packageDescription": string,
  "packageDestination": string,
  "packageDays": integer,
  "packageNights": integer,
  "packageAccommodation": string,
  "packageTransportation": string,
  "packageMeals": string,
  "packageActivities": string,
  "packagePrice": float,
  "packageDiscountPrice": float,
  "packageOffer": boolean,
  "packageRating": 0,
  "packageTotalRatings": 0
}}

TRAVEL PLAN:
\"\"\"
{data.plan_text}
\"\"\"
"""

    try:
        raw_response = generate_with_gemini(extraction_prompt)
        print("Gemini raw response:", raw_response)

        # Extract the JSON body from the response
        json_match = re.search(r'\{.*\}', raw_response, re.DOTALL)
        if not json_match:
            raise ValueError("No JSON object found in Gemini response")

        json_str = json_match.group(0)
        structured_package = json.loads(json_str)

        # Add image results based on the extracted destination
        destination = structured_package.get("packageDestination", "Suggested Destination")
        image_queries = create_image_queries({}, destination)
        images = image_searcher.search_all_sources(image_queries, num_images=6)

        structured_package["packageImages"] = images

    except Exception as e:
        print("Gemini extraction failed:", e)
        print("Prompt was:\n", extraction_prompt)
        raise HTTPException(status_code=500, detail="Failed to extract structured package data.")

    return structured_package



class ChatRequest(BaseModel):
    plan_text: str
    user_query: str

def create_gemini_prompt(plan_text: str, user_query: str) -> str:
    return f"""
You are a helpful travel assistant specialized in tours and travel within Pakistan. 

Here is a travel plan describing a trip:

\"\"\"
{plan_text}
\"\"\"

Please answer the user's question below **only if it is related to travel, sightseeing, food, transport, costs, culture, or activities in Pakistan based on the above plan**.

User question:
{user_query}

If the question is unrelated to travel in Pakistan, respond politely with:
"I can only answer questions related to travel and tours within Pakistan."

Answer concisely and clearly.
"""

@app.post("/chat", summary="Answer user questions about the generated travel plan")
async def chat_about_plan(chat_input: ChatRequest):
    try:
        prompt = create_gemini_prompt(chat_input.plan_text, chat_input.user_query)
        llm_response = generate_with_gemini(prompt)  # Remove 'await' if your wrapper is not async
        return {"response": llm_response.strip()}
    except Exception as e:
        print("Error in chat endpoint:", e)
        raise HTTPException(status_code=500, detail="Failed to process your question.")