# Standard library imports
import os
import re
import sqlite3
from typing import List, Tuple, Optional

# Third-party imports
import gradio as gr
from dotenv import load_dotenv

# SmolaGents imports
from smolagents import CodeAgent, LiteLLMModel, tool
from smolagents.agent_types import AgentText
from smolagents.memory import ActionStep, TaskStep
# Local imports
from init_db import Project

# Load environment variables
load_dotenv()

@tool
def sql_engine(query: str) -> List[Tuple]:
    """
    Allows you to perform SQL queries on the table in SQLite database. Returns list of results.
    The table is named 'project'. Its description is as follows:
        Columns:
        - url: VARCHAR(255) - URL of the project
        - title: VARCHAR(255) - Title of the project
        - deposit: INTEGER - Percentage of payment before completion
        - min_price: INTEGER - Lowest available apartment price in CZK with VAT
        - status: VARCHAR(255) - Status of the project (preparation, selling, sold out)
        - city: VARCHAR(255) - City of the project
        - lat: FLOAT - GPS latitude coordinates
        - lng: FLOAT - GPS longitude coordinates
        - start_year: INTEGER - Year of construction start
        - end_year: INTEGER - Year of construction end
        - developer: VARCHAR(255) - Name of the construction company
        - ignore: BOOLEAN - If True, the project does not have any apartments for sale
        
        # Apartment prices by type (all INTEGER in CZK with VAT)
        - price_1kk: Price of 1+kk apartment
        - price_2kk: Price of 2+kk apartment
        - price_3kk: Price of 3+kk apartment
        - price_4kk: Price of 4+kk apartment
        - price_5kk: Price of 5+kk apartment
        - price_6kk: Price of 6+kk apartment
        - price_7kk: Price of 7+kk apartment
        - price_8kk: Price of 8+kk apartment
        - price_9kk: Price of 9+kk apartment
        - price_10kk: Price of 10+kk apartment
        
        - price_1_1: Price of 1+1 apartment
        - price_2_1: Price of 2+1 apartment
        - price_3_1: Price of 3+1 apartment
        - price_4_1: Price of 4+1 apartment
        - price_5_1: Price of 5+1 apartment
        - price_6_1: Price of 6+1 apartment
        - price_7_1: Price of 7+1 apartment
        - price_8_1: Price of 8+1 apartment
        - price_9_1: Price of 9+1 apartment
        - price_10_1: Price of 10+1 apartment
        
        - content: TEXT - Raw content (ignore to save tokens)
        - created_at: DATETIME - Date and time of creation

    Args:
        query: The query to perform. This should be correct SQL.
    """
    con = sqlite3.connect("estate.db")

    result = con.execute(query)
    all_rows = result.fetchall()

    con.close()
    return all_rows



sql_distance_query = """
SELECT 
    id, 
    url,
    structure->>'$.lat' AS lat,
    structure->>'$.lng' AS lng,
    (6371000 * acos(
        cos(radians(50.08804)) * cos(radians(CAST(structure->>'$.lat' AS FLOAT))) * 
        cos(radians(CAST(structure->>'$.lng' AS FLOAT)) - radians(14.42076)) + 
        sin(radians(50.08804)) * sin(radians(CAST(structure->>'$.lat' AS FLOAT)))
    )) AS distance
FROM project
WHERE structure->>'$.lat' IS NOT NULL 
AND structure->>'$.lng' IS NOT NULL
ORDER BY distance ASC
LIMIT 5;
"""


litemodel = LiteLLMModel("o3-mini-2025-01-31", api_base="https://api.openai.com/v1", api_key=os.getenv("OPENAI_API_KEY"))
codeagent = CodeAgent(tools=[sql_engine], model=litemodel)

def chat_with_sql(message, history):
    try:
      input_message = "You are a helpful assistant that presents information about real estate projects in database for investors. Use 'IS NOT NULL' when necessary. Translate any city names from input into Czech language. Present useful fields such as url. Do not make up information or hallucinate. If final result rows are empty respond that you could not find records matching criteria. Use only Czech city names. Note that there is limit 10 records for database result. Reformat final answer in markdown. User Input:\n\n {message}"
      prompt = input_message.format(message=message)
      result = codeagent.run(prompt, reset=False)
      yield  gr.ChatMessage(role="assistant", content=result)

    except Exception as e:
        yield gr.ChatMessage(role="assistant", content=f"Error: {str(e)}")

# Create the Gradio interface
count_projects = Project.select().where(Project.structure.is_null(False)).count()
demo = gr.ChatInterface(
    fn=chat_with_sql,
    title="Estate Chat",
    description=f"Planning to buy a new apartment? Ask me about {count_projects} new constructions (novostavby) in Prague 🇨🇿",
    examples=[
        "Cheapest apartment in Prague",
        "Project with the lowest deposit before completion",
        "All 2+kk below 8000000 CZK",
    ],
    type="messages"
)

if __name__ == "__main__":
    demo.launch(share=False, debug=True)