from fastapi import FastAPI, HTTPException
from smolagents import CodeAgent, HfApiModel, tool, DuckDuckGoSearchTool
from pydantic import BaseModel
from typing import List
import re
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

app = FastAPI()

class Query(BaseModel):
    text: str

class SearchResponse(BaseModel):
    urls: List[str]
    full_response: str

@tool
def generate_bible_url(book_num: str, chapter: str, verse: str) -> str:
    """Generates a URL for a bible verse using the JW.org format

    Args:
        book_num: two digit book number (01-66)
        chapter: three digit chapter number (padded with zeros)
        verse: three digit verse number (padded with zeros)
    """
    bible_code = f"{book_num.zfill(2)}{chapter.zfill(3)}{verse.zfill(3)}"
    return f"https://www.jw.org/finder?srcid=jwlshare&wtlocale=E&prefer=lang&bible={bible_code}&pub=nwtsty"

# Initialize the agent with Qwen model and token
agent = CodeAgent(
    tools=[generate_bible_url, DuckDuckGoSearchTool()],
    model=HfApiModel(
        "Qwen/Qwen2.5-Coder-32B-Instruct",
        token=os.getenv("HUGGINGFACE_API_TOKEN")
    ),
    additional_authorized_imports=[]
)

@app.post("/search", response_model=SearchResponse)
async def search_verses(query: Query):
    try:
        # Execute the agent with the query
        result = agent.run(f"""You are a helpful AI assistant specializing in finding relevant Bible verses.

        1. First use the DuckDuckGo search tool to search specifically for "{query.text} bible verse"
        2. From the search results, identify the most relevant verse that directly addresses the query
        3. Verify the verse's relevance before generating the URL
        4. Use the generate_bible_url tool to create the URL for the chosen verse

        Only return verses that specifically address: {query.text}

        Query: {query.text}""")

        # Extract URLs whether result is string or list
        if isinstance(result, list):
            # Handle list of tuples (url, description)
            urls = [item[0] for item in result if isinstance(item, tuple) and len(item) > 0]
        else:
            # Handle string response
            urls = re.findall(r'https://www\.jw\.org/finder\?[^\s\)]+', str(result))

        return {
            "urls": urls,
            "full_response": str(result)
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)