from fasthtml_hf import setup_hf_backup
from fasthtml.common import (
    picolink,
    serve,
    Div,
    Title,
    Main,
    Input,
    Button,
    A,
    Section,
    H2,
    Ul,
    Li,
    P,
    Img,
    Details,
    MarkdownJS,
    HighlightJS,
    Summary,
    Script,
    I,
    Form,
    RedirectResponse,
    dataclass,
    Favicon,
    database,
    get_key,
    Table,
    Thead,
    Tr,
    Th,
    Tbody,
    Td,
    FileResponse,
    fast_app,
    Beforeware,
    Hidden,
    Request,
    H3,
    Style,
)
from fasthtml.components import Nav, Article, Header, Mark
from fasthtml.pico import Search, Grid, Fieldset, Label
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.sessions import SessionMiddleware
from starlette.middleware.cors import CORSMiddleware
from starlette.middleware.authentication import AuthenticationMiddleware
from starlette.authentication import AuthenticationBackend, AuthenticationError, SimpleUser, AuthCredentials
from starlette.responses import RedirectResponse
from vespa.application import Vespa
import json
import os
import re
import time
from hmac import compare_digest
from io import StringIO
import csv
import tempfile
from enum import Enum
from typing import Tuple as T
from urllib.parse import quote
import uuid
import secrets

DEV_MODE = False

if DEV_MODE:
    print("Running in DEV_MODE - Hot reload enabled")
    print("Loading environment variables from .env")
    from dotenv import load_dotenv

    load_dotenv()
else:
    print("DEV_MODE disabled - environment variables loaded from system")

vespa_app_url = os.getenv("VESPA_APP_URL", None)
if vespa_app_url is None:
    print("Please set the VESPA_APP_URL environment variable")
    exit(1)

ADMIN_NAME = os.getenv("ADMIN_NAME", "admin")
ADMIN_PWD = os.getenv("ADMIN_PWD", "admin")

vespa_app: Vespa = Vespa(
    url=vespa_app_url,
    vespa_cloud_secret_token=os.getenv("VESPA_CLOUD_SECRET_TOKEN"),
)
status = vespa_app.get_application_status()
if status is None:
    print("Could not connect to Vespa application")
else:
    print("Connected to Vespa application!")

fa = Script(src="https://kit.fontawesome.com/664eb1a115.js", crossorigin="anonymous")
favicon = Favicon(
    "https://search.vespa.ai/favicon.ico",
    "https://search.vespa.ai/favicon.ico",
)
DB_FILE = "db/vespa.db"
db = database(DB_FILE)
queries = db.t.queries
if queries not in db.t:
    # You can pass a dict, or kwargs, to most MiniDataAPI methods.
    queries.create(
        dict(qid=int, query=str, ranking=str, sess_id=str, timestamp=int), pk="qid"
    )
    # Add autoincrement to the qid column
    db.query("ALTER TABLE queries ADD COLUMN qid INTEGER PRIMARY KEY AUTOINCREMENT")
Query = queries.dataclass()

# Add a classmethod to the Query dataclass to convert timestamp field to a human readable format
Query.get_datetime = lambda self: time.strftime(
    "%Y-%m-%d %H:%M:%S", time.localtime(self.timestamp)
)

# Status code 303 is a redirect that can change POST to GET,
# so it's appropriate for a login page.
login_redir = RedirectResponse("/login", status_code=303)


def user_auth_before(req, sess):
    sess.setdefault('auth', False)


spinner_css = Style("""
    .htmx-indicator {
        display: none; /* Hide spinner by default */
    }

    .htmx-indicator.htmx-request {
        display: block;    
    }
""")

headers = (
    picolink,
    MarkdownJS(),
    HighlightJS(langs=["json", "python"]),
    favicon,
    fa,
    spinner_css,
)

# Read file contents once before starting the server
with open("README.md") as f:
    README = f.read()
with open("main.py") as f:
    SOURCE = f.read()

# Sesskey
sess_key_path = "session/.sesskey"
SESSION_KEY = "session_"
# Make sure session directory exists
os.makedirs("session", exist_ok=True)


# Middleware
class XFrameOptionsMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        response = await call_next(request)
        response.headers["X-Frame-Options"] = "ALLOW-FROM https://huggingface.co/"
        return response

class SessionLoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request, call_next):
        print(f"Before request: Session data: {request.session}")
        response = await call_next(request)
        print(f"After request: Session data: {request.session}")
        return response

class DebugSessionMiddleware(SessionMiddleware):
    async def __call__(self, scope, receive, send):
        print(f"DebugSessionMiddleware: Before processing - Scope: {scope}")
        await super().__call__(scope, receive, send)
        print(f"DebugSessionMiddleware: After processing - Scope: {scope}")

# Generate a secure secret key
SECRET_KEY = secrets.token_urlsafe(32)

# Modify the middleware setup
middlewares = [
    Middleware(
        SessionMiddleware,
        secret_key=SECRET_KEY,
        max_age=3600,
        same_site='None',  # Try 'Lax' if 'None' doesn't work
        https_only=True,  # Set to True if your site uses HTTPS
    ),
]

bware = Beforeware(
    user_auth_before,
    skip=[
        r"/favicon\.ico",
        r"/static/.*",
        r".*\.css",
        r".*\.js",
    ],
)

app, rt = fast_app(
    before=bware,
    live=DEV_MODE,
    hdrs=headers,
    #middleware=middlewares,
    key_fname=sess_key_path,
    samesite="none",
    secure=True,  # Add this line
    httponly=True,  # Add this line
)

# Add this function for debugging
@app.middleware("http")
async def debug_request_middleware(request: Request, call_next):
    print(f"Request headers: {request.headers}")
    print(f"Request cookies: {request.cookies}")
    response = await call_next(request)
    print(f"Response headers: {response.headers}")
    if hasattr(response, 'cookies'):
        print(f"Response cookies: {response.cookies}")
    return response

sesskey = get_key(fname=sess_key_path)
print(f"Session key: {sesskey}")


# enum class for rank profiles
class RankProfile(str, Enum):
    bm25 = "bm25"
    semantic = "semantic"
    fusion = "fusion"


def get_navbar(admin: bool):
    print(f"In get_navbar: {admin}")
    bar = Nav(
        Ul(
            Li(
                A(
                    Img(src="https://vespa.ai/assets/vespa-ai-logo-heather.svg"),
                    href="https://cloud.vespa.ai",
                    target="_blank",
                    style="margin: 10px;",
                ),
            )
        ),
        Ul(H2("Vespa-fastHTML demo")),
        Ul(
            # A question mark icon with link to an about page
            A(
                I(cls="fa fa-question-circle fa-2x"),
                href="/about",
                style="margin: 10px;",
                title="About this app",
            ),
            A(
                I(cls="fab fa-slack fa-2x"),
                href="https://slack.vespa.ai/",
                style="margin: 10px;",
                target="_blank",
                title="Join Vespa Slack channel",
            ),
            A(
                I(cls="fab fa-github fa-2x"),
                href="https://github.com/vespa-engine/sample-apps/tree/master/examples/fasthtml-demo",
                style="margin: 10px;",
                target="_blank",
                title="View source code on GitHub",
            ),
            A(
                I(cls="fa fa-code fa-2x"),
                href="/source",
                style="margin: 10px;",
                title="View source code",
            ),
            # Login icon (link to /login) show tooltip on hover. MAke it hidden if admin is logged in
            A(
                I(cls="fa fa-shield fa-2x"),
                href="/login" if not admin else "/admin",
                style="margin: 10px;",
                title="Admin login",
            ),
            # Logout icon if admin is logged in
            A(
                I(cls="fa fa-sign-out fa-2x"),
                href="/logout",
                style="margin: 10px;" if admin else "display: none;",
                title="Logout",
            ),
        ),
        # 10px margin to right of navbar
        style="margin-right: 10px;",
    )
    return bar


def spinner_div(hidden: bool = False):
    return Div(
        A(
            id="spinner",
            aria_busy="true",
            cls="htmx-indicator",
            style="font-size: 2em;",
        ),
        style="text-align: center; margin-top: 40px;"
        if not hidden
        else "display: none;",
    )


@app.route("/")
def get(sess):
    auth = sess.get('auth', False)
    queries = [
        "Breast Cancer Cells Feed on Cholesterol",
        "Treating Asthma With Plants vs. Pills",
        "Testing Turmeric on Smokers",
        "The Role of Pesticides in Parkinson's Disease",
    ]
    return (
        Title("Vespa demo"),
        get_navbar(auth),
        Main(
            # Search bar
            Search(
                Input(
                    type="search",
                    placeholder="Ask/search for medical information?",
                    id="userquery",
                ),
                # Get search results on button click with search-input as query parameter
                Button(
                    "Search",
                    hx_get="/search",
                    # include userquery and id of selected ranking radio button
                    hx_include="#userquery, input[name=ranking]:checked",
                    hx_target="#results",
                    hx_indicator="#spinner",
                ),
                style="margin: 10% 10px 0 0;",
            ),
            Fieldset(
                Input(type="radio", id="bm25", name="ranking", value="bm25"),
                Label("BM25", htmlfor="bm25"),
                Input(type="radio", id="semantic", name="ranking", value="semantic"),
                Label("Semantic", htmlfor="semantic"),
                Input(
                    type="radio",
                    id="fusion",
                    name="ranking",
                    value="fusion",
                    checked="",
                ),
                Label("Reciprocal Rank fusion", htmlfor="fusion"),
                style="margin: 10px; text-align: center;",
                id="ranking",
            ),
            H3("Example queries"),
            # Buttons with predefined search queries
            Grid(
                *[
                    Button(
                        query,
                        hx_get="/search?userquery=" + query,
                        hx_include="input[name=ranking]:checked",
                        hx_target="#results",
                        hx_indicator="#spinner",
                        hx_on_click=f"document.getElementById('userquery').value='{query}'",
                        style="margin: 10px; padding: 5px;",
                        cls="secondary outline",
                        id=f"example-{qid}",
                    )
                    for qid, query in enumerate(queries)
                ],
                # Make the grid buttons have same height and distribute evenly and center align
                style="grid-template-columns: repeat(auto-fill, minmax(200px, 1fr));",
            ),
            # Section(
            #     Input(
            #         id="suggestion-input",
            #         list="search-options",
            #         placeholder="Search options",
            #     ),
            #     Datalist(
            #             *[
            #                 Option(
            #                     "Covid-19",
            #                     value="Covid-19",
            #                 ),
            #                 Option(
            #                     "Vaccine",
            #                     value="Vaccine",
            #                 ),
            #             ],
            #             id="search-options",
            #         ),
            #     id="suggestions",
            # ),
            # Display spinner div only if it #spinner does not exist
            Section(
                spinner_div(),
                id="results",
                hx_swap="innerHTML",
                style="margin: 20px;",
            ),
            style="margin: 0 auto; width: 70%;",
            id="main",
        ),
    )


@dataclass
class Login:
    name: str
    pwd: str


@app.get("/login")
def get_login_form(sess, error: bool = False):
    auth = sess.get('auth', False)
    frm = Form(
        Input(id="name", placeholder="Name"),
        Input(id="pwd", type="password", placeholder="Password"),
        Button("login"),
        action="/login",
        method="post",
    )
    err_msg = P("Incorrect password", style="color: red;") if error else ""
    return (
        Title("Admin login"),
        get_navbar(auth),
        Main(
            err_msg,
            frm,
            style="width: 50%; margin: 10% auto;",
        ),
    )


@app.post("/login")
async def login(sess, request: Request):
    form = await request.form()
    username = form.get("name")
    password = form.get("pwd")
    
    if username == ADMIN_NAME and compare_digest(ADMIN_PWD.encode("utf-8"), password.encode("utf-8")):
        sess[SESSION_KEY] = {'auth': True}
        return RedirectResponse("/admin", status_code=303)
    
    return RedirectResponse("/login?error=True", status_code=303)


@app.route("/logout")
async def logout(sess):
    if SESSION_KEY in sess:
        del sess[SESSION_KEY]
    return RedirectResponse("/")


def replace_hi_with_strong(text):
    parts = re.split(r"(<hi>|</hi>)", text)
    elements = []
    open_tag = False
    for part in parts:
        if part == "<hi>":
            open_tag = True
        elif part == "</hi>":
            open_tag = False
        elif open_tag:
            elements.append(Mark(part))
        else:
            elements.append(part)
    return elements


def log_query_to_db(query, ranking, sess):
    queries.insert(
        Query(query=query, ranking=ranking, sess_id=sesskey, timestamp=int(time.time()))
    )
    if 'user_id' not in sess:
        sess['user_id'] = str(uuid.uuid4())
    
    sess.setdefault('queries', [])
    
    query_data = {
        'query': query,
        'ranking': ranking,
        'timestamp': int(time.time())
    }
    sess['queries'].append(query_data)
    
    # Limit the number of queries stored in the session to prevent it from growing too large
    sess['queries'] = sess['queries'][-100:]  # Keep only the last 100 queries

    return query_data


def parse_results(records):
    return [
        Article(
            Header(
                H2(
                    A(
                        result["title"],
                        hx_get=f"/document/{result['id']}",
                        hx_target="#results",
                    )
                )
            ),
            Div(
                P(
                    *replace_hi_with_strong(
                        result["body"][:300] + "..."
                    ),  # Display first 300 characters of body
                ),
                Div(
                    # Button with "Show more" - center align
                    Button(
                        "Show more",
                        hx_post=f"/expand/{result['id']}?expand=true",
                        hx_target=f"#{result['id']}",
                        hx_include=f"#{result['id']}-full",
                        cls="outline secondary",
                        # Style to fill whole width of parent div
                        style="width: 100%;",
                    ),
                    style="text-align: center;",
                ),
                id=result["id"],
            ),
            Hidden(result["body"], id=f"{result['id']}-full"),
        )
        for result in records
    ]


@app.post("/expand/{docid}")
async def expand(request: Request, docid: str, expand: bool):
    print(f"Expanding {docid}")
    form_data = await request.form()
    result = form_data.get(f"{docid}-full")
    if not expand:
        result = result[:300] + "..."
    return (
        Div(
            P(
                *replace_hi_with_strong(result),  # Display full body
            ),
            Div(
                # Button with "Show less" - center align
                Button(
                    "Show less" if expand else "Show more",
                    hx_post=f"/expand/{docid}?expand="
                    + ("false" if expand else "true"),
                    hx_target=f"#{docid}",
                    hx_include=f"#{docid}-full",
                    cls="outline secondary",
                    # Style to fill whole width of parent div
                    style="width: 100%;",
                ),
                style="text-align: center;",
            ),
            id=docid,
        ),
    )


# Returns tuple of (yql, body(dict)) based on the ranking profile
def get_yql(ranking: RankProfile, userquery: str) -> T[str, dict]:
    if ranking == RankProfile.bm25:
        yql = "select * from sources * where userQuery() limit 10"
        body = {}
    elif ranking == RankProfile.semantic:
        yql = "select * from sources * where ({targetHits:10}nearestNeighbor(embedding,q)) limit 10"
        body = {"input.query(q)": f"embed({userquery})"}
    elif ranking == RankProfile.fusion:
        yql = "select * from sources * where rank({targetHits:1000}nearestNeighbor(embedding,q), userQuery()) limit 10"
        body = {"input.query(q)": f"embed({userquery})"}
    return yql, body


@app.get("/search")
async def search(sess, userquery: str, ranking: str):
    print(sess)
    quoted = quote(userquery) + "&ranking=" + ranking
    log_query_to_db(userquery, ranking, sess)
    yql, body = get_yql(ranking, userquery)
    async with vespa_app.asyncio() as session:
        resp = await session.query(
            yql=yql,
            query=userquery,
            hits=10,
            ranking=str(ranking),
            body=body,
        )
    records = []
    fields = ["id", "title", "body"]
    for hit in resp.hits:
        record = {}
        for field in fields:
            record[field] = hit["fields"][field]
        records.append(record)
    results = parse_results(records)
    json_dump = json.dumps(resp.get_json(), indent=4)
    return Div(
        spinner_div(),
        # Accordion (with Details)
        Details(
            Summary("Full JSON response"),
            Div(
                f"""```json\n{json_dump}\n```""",
                cls="marked",
            ),
        ),
        H2(
            "Search Results",
        ),
        Div(
            *results,
            id="all-searchresults",
        ),
    )


@app.get("/download_csv")
def download_csv(request: Request):
    queries_dict = list(db.query("SELECT * FROM queries"))
    queries = [Query(**query) for query in queries_dict]

    # Create CSV in memory
    csv_file = StringIO()
    csv_writer = csv.writer(csv_file)
    csv_writer.writerow(["Query", "Session ID", "Timestamp"])
    for query in queries:
        csv_writer.writerow([query.query, query.sess_id, query.timestamp])

    # Move to the beginning of the StringIO object
    csv_file.seek(0)

    # Save CSV to a temporary file
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
    temp_file.write(csv_file.getvalue().encode("utf-8"))
    temp_file.close()

    return FileResponse(
        temp_file.name,
        filename="queries.csv",
        media_type="text/csv",
        content_disposition_type="attachment",
    )


@app.route("/admin")
async def admin(request, sess):
    auth = sess.get(SESSION_KEY, {}).get('auth', False)
    if not auth:
        print(f"Not authenticated: {auth}")
        return RedirectResponse("/login", status_code=303)
    
    page = int(request.query_params.get("page", 1))
    limit = 15
    offset = (page - 1) * limit
    total_queries_result = list(
        db.query("SELECT COUNT(*) AS count FROM queries ORDER BY timestamp DESC")
    )
    total_queries = total_queries_result[0]["count"]
    queries_dict = list(
        db.query(f"SELECT * FROM queries LIMIT {limit} OFFSET {offset}")
    )
    queries = [Query(**query) for query in queries_dict]

    total_pages = (
        total_queries + limit - 1
    ) // limit  # Calculate total number of pages

    # Define the range of pages to display
    page_window = 5  # Number of pages to display at once
    start_page = max(1, page - page_window // 2)
    end_page = min(total_pages, start_page + page_window - 1)

    # Adjust the start and end pages if they exceed the limits
    if end_page - start_page < page_window:
        start_page = max(1, end_page - page_window + 1)

    # Pagination controls with "First", "Previous", "Next", and "Last"
    pagination_controls = Div(
        A(
            "First",
            href="/admin?page=1",
            style="margin: 5px;"
            if page > 1
            else "margin: 5px; color: grey; pointer-events: none;",
        ),
        A(
            "Previous",
            href=f"/admin?page={page - 1}",
            style="margin: 5px;"
            if page > 1
            else "margin: 5px; color: grey; pointer-events: none;",
        ),
        *[
            A(
                f"{i}",
                href=f"/admin?page={i}",
                style="margin: 5px;"
                if i != page
                else "margin: 5px; font-weight: bold;",
            )
            for i in range(start_page, end_page + 1)
        ],
        A(
            "Next",
            href=f"/admin?page={page + 1}",
            style="margin: 5px;"
            if page < total_pages
            else "margin: 5px; color: grey; pointer-events: none;",
        ),
        A(
            "Last",
            href=f"/admin?page={total_pages}",
            style="margin: 5px;"
            if page < total_pages
            else "margin: 5px; color: grey; pointer-events: none;",
        ),
        style="text-align: center; margin: 20px;",
    )

    # Total pages indication
    total_pages_indicator = Div(
        f"Page {page} of {total_pages}",
        style="text-align: center; margin: 10px;",
    )

    return (
        Title("Admin"),
        get_navbar(auth),
        Main(
            Div(
                A(
                    I(cls="fa fa-arrow-left"),
                    "Back",
                    href="/",
                    title="Back to main page",
                    style="margin: 10px;",
                ),
                style="margin: 10px;",
            ),
            H2("Queries"),
            # Table of all queries
            Table(
                Thead(
                    Tr(
                        Th("Query"),
                        Th("Session ID"),
                        Th("Datetime"),
                    )
                ),
                Tbody(
                    *[
                        Tr(
                            Td(query.query),
                            Td(query.sess_id),
                            Td(query.get_datetime()),
                        )
                        for query in queries
                    ],
                ),
                cls="striped",
            ),
            total_pages_indicator,  # Include the total pages indicator here
            pagination_controls,
            Div(
                A(
                    I(cls="fa fa-download fa-2x"),
                    " Download CSV",
                    href="/download_csv",
                    style="margin: 10px; float: right;",
                    title="Download queries as CSV",
                ),
                style="text-align: right; margin: 20px;",
            ),
            style="width: 80%; margin: 40px auto;",
        ),
    )


@app.get("/source")
def get_source(auth, sess):
    # Back icon to go back to main page in top left corner
    return (
        Title("Source code"),
        get_navbar(auth),
        Main(
            Div(
                A(
                    I(cls="fa fa-arrow-left"),
                    "Back",
                    href="/",
                    title="Back to main page",
                    style="margin: 10px;",
                ),
                Div(
                    f"""### `main.py`\n### This is the complete source code for this app \n```python\n{SOURCE}\n```""",
                    cls="marked",
                    style="margin: 10px;",
                ),
                style="width: 80%; margin: 40px auto;",
            ),
        ),
    )


@app.get("/about")
def get_about(auth, sess):
    # Strip everything before the FIRST # in the README
    stripped_readme = re.sub(
        r"^.*?(?=# FastHTML Vespa frontend)", "", README, flags=re.DOTALL
    )

    return (
        Title("About this app"),
        get_navbar(auth),
        Main(
            Div(
                A(
                    I(cls="fa fa-arrow-left"),
                    "Back",
                    href="/",
                    title="Back to main page",
                    style="margin: 10px;",
                ),
                Div(
                    stripped_readme,
                    cls="marked",
                    style="margin: 10px;",
                ),
                style="width: 80%; margin: 40px auto;",
            ),
        ),
    )


@app.get("/document/{docid}")
def get_document(docid: str, sess):
    resp = vespa_app.get_data(data_id=docid, schema="doc", namespace="tutorial")
    doc = resp.json
    # Link with Back to search results at top of page
    last_query = sess.get('queries', [{}])[-1].get('query', '')
    return Main(
        Div(
            A(
                I(cls="fa fa-arrow-left"),
                "Back to search results",
                hx_get=f"/search?userquery={last_query}",
                hx_target="#results",
                style="margin: 10px;",
            ),
            H2(doc["fields"]["title"], style="margin: 10px;"),
            P(doc["fields"]["body"], cls="marked"),
        ),
    )


if not DEV_MODE:
    try:
        setup_hf_backup(app)
    except Exception as e:
        print(f"Error setting up hf backup: {e}")
serve()