import os
import gradio as gr
from smolagents import CodeAgent, tool, LiteLLMModel
from qdrant_client import QdrantClient
from qdrant_client.models import Filter, FieldCondition, MatchValue, PointStruct
from sentence_transformers import SentenceTransformer
from firecrawl import FirecrawlApp

# --- Setup ---
client = QdrantClient(url=os.getenv("QDRANT_URL"), api_key=os.getenv("QDRANT_API_KEY"), check_compatibility=False)
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
firecrawl = FirecrawlApp(api_key=os.getenv("FIRECRAWL_API_KEY"))
collection_name = "cricket_analysis"

# --- Tools ---
@tool
def search_cricket_database(query: str, active_match_id: str) -> str:
    """
    Search the ball-by-ball database for match events.
    Args:
        query: The search term about the match.
        active_match_id: The ID to filter results.
    """
    query_vector = embed_model.encode(query).tolist()
    try:
        # Try modern query method first
        search_result = client.query_points(
            collection_name=collection_name,
            query=query_vector,
            query_filter=Filter(must=[FieldCondition(key="match_id", match=MatchValue(value=active_match_id))]),
            limit=5
        ).points
    except AttributeError:
        # Fallback for different client versions
        search_result = client.search(
            collection_name=collection_name,
            query_vector=query_vector,
            query_filter=Filter(must=[FieldCondition(key="match_id", match=MatchValue(value=active_match_id))]),
            limit=5
        )
    
    if not search_result: return "No data found."
    return "\n\n".join([f"[Context] {r.payload.get('text')}" for r in search_result])

# --- Agent Initialization ---
model = LiteLLMModel(model_id="huggingface/Qwen/Qwen2.5-72B-Instruct", api_key=os.getenv("HF_TOKEN_cric"))

agent = CodeAgent(
    tools=[search_cricket_database], 
    model=model,
    # 🛡️ Allowing the agent to use specific imports if it generates code using them
    additional_authorized_imports=["web_search", "math", "time", "re"]
)

# --- UI Logic ---
def chat_fn(message, history, mode, custom_url, custom_id):
    active_id = custom_id if mode == "Custom Match" else "T20_WC_Final_2024"
    
    # Run the agent
    response = agent.run(f"Match ID: {active_id}. Question: {message}")
    
    # ✅ History update using dictionary format
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})
    
    return "", history

with gr.Blocks() as demo:
    gr.Markdown("# 🏏 Cricket Analytics Portal")
    mode_select = gr.Radio(["Featured Match", "Custom Match"], label="Mode", value="Featured Match")
    chatbot = gr.Chatbot(label="Analysis Feed", type="messages")
    msg = gr.Textbox(label="Ask the Analyst")
    msg.submit(chat_fn, [msg, chatbot, mode_select, gr.State(""), gr.State("")], [msg, chatbot])

if __name__ == "__main__":
    demo.launch()