from __future__ import annotations

import hashlib
import hmac
import json
import os
import threading
from collections import Counter
from datetime import timezone, date, datetime, timedelta
from pathlib import Path
from typing import Annotated, Any

import gradio as gr
try:
    from geoip2.database import Reader
    from geoip2.errors import AddressNotFoundError
except ImportError:
    Reader = None

    class AddressNotFoundError(Exception):
        pass
#import spaces
#import torch

#try:
    #import spaces
#except ImportError:
    ## The `spaces` package is injected by the Hugging Face ZeroGPU runtime.
    ## Keep local development and CPU-based tests usable without that runtime.
    #class _SpacesFallback:
        #@staticmethod
        #def GPU(*_args, **_kwargs):
            #def decorator(function):
                #return function

            #return decorator

    #spaces = _SpacesFallback()

from fastapi import FastAPI, Header, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response
from pydantic import BaseModel, ConfigDict, Field, field_validator

APP_TITLE = "HF Space Analytics"
DATA_DIR = Path(os.getenv("DATA_DIR", "/data/visits"))
TRACKING_KEY = os.getenv("TRACKING_KEY", "")
ADMIN_KEY = os.getenv("ADMIN_KEY", "")
HASH_SALT = os.getenv("HASH_SALT", "change-this-salt")
GEOIP_DB_PATH = Path(os.getenv("GEOIP_DB_PATH", "/data/GeoLite2-City.mmdb"))
MAX_FIELD_LENGTH = 500
WRITE_LOCK = threading.Lock()
GEOIP_READER: Any = None

if Reader is not None and GEOIP_DB_PATH.is_file():
    try:
        GEOIP_READER = Reader(str(GEOIP_DB_PATH))
    except OSError:
        GEOIP_READER = None

DATA_DIR.mkdir(parents=True, exist_ok=True)

origins_raw = os.getenv("ALLOWED_ORIGINS", "*")
ALLOWED_ORIGINS = [x.strip() for x in origins_raw.split(",") if x.strip()]

api = FastAPI(title=APP_TITLE, version="2.0.0")
api.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=False,
    allow_methods=["GET", "POST", "OPTIONS"],
    allow_headers=["Content-Type", "X-Tracking-Key", "X-Admin-Key"],
)


class TrackEvent(BaseModel):
    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)

    space: str = Field(min_length=1, max_length=200)
    path: str | None = Field(default=None, max_length=MAX_FIELD_LENGTH)
    referrer: str | None = Field(default=None, max_length=MAX_FIELD_LENGTH)
    visitor_id: str | None = Field(default=None, max_length=200)
    timezone: str | None = Field(default=None, max_length=100)
    language: str | None = Field(default=None, max_length=100)
    event: str = Field(default="page_view", min_length=1, max_length=100)

    @field_validator("space")
    @classmethod
    def validate_space(cls, value: str) -> str:
        if any(c in value for c in "\r\n\t"):
            raise ValueError("space contains invalid control characters")
        return value


def secure_equal(provided: str | None, expected: str) -> bool:
    return bool(provided) and hmac.compare_digest(provided, expected)


def require_tracking_key(key: str | None) -> None:
    if TRACKING_KEY and not secure_equal(key, TRACKING_KEY):
        raise HTTPException(status_code=401, detail="Invalid tracking key")


def require_admin_key(key: str | None) -> None:
    if ADMIN_KEY and not secure_equal(key, ADMIN_KEY):
        raise HTTPException(status_code=401, detail="Invalid admin key")


def hash_visitor(visitor_id: str | None) -> str | None:
    if not visitor_id:
        return None
    digest = hashlib.sha256(f"{HASH_SALT}:{visitor_id}".encode()).hexdigest()
    return digest[:32]


def hash_client_ip(ip_address: str | None) -> str | None:
    if not ip_address:
        return None
    digest = hashlib.sha256(f"{HASH_SALT}:ip:{ip_address}".encode()).hexdigest()
    #return digest[:32]
    return ip_address


def lookup_location(ip_address: str | None) -> dict[str, str] | None:
    if not ip_address or GEOIP_READER is None:
        return None
    try:
        response = GEOIP_READER.city(ip_address)
    except (AddressNotFoundError, ValueError):
        return None

    location = {
        "country_code": response.country.iso_code,
        "country": response.country.name,
        "region": response.subdivisions.most_specific.name,
        "city": response.city.name,
    }
    return {key: value for key, value in location.items() if value}


def daily_file(day: date) -> Path:
    return DATA_DIR / f"{day.isoformat()}.jsonl"


def append_event(record: dict[str, Any]) -> None:
    path = daily_file(datetime.now(timezone.utc).date())
    line = json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n"
    with WRITE_LOCK:
        with path.open("a", encoding="utf-8") as handle:
            handle.write(line)
            handle.flush()
            os.fsync(handle.fileno())


def iter_events(days: int):
    today = datetime.now(timezone.utc).date()
    for offset in range(days - 1, -1, -1):
        path = daily_file(today - timedelta(days=offset))
        if not path.exists():
            continue
        try:
            with path.open("r", encoding="utf-8") as handle:
                for line in handle:
                    try:
                        yield json.loads(line)
                    except json.JSONDecodeError:
                        continue
        except OSError:
            continue


def calculate_stats(days: int, space: str | None = None) -> dict[str, Any]:
    events = [
        row
        for row in iter_events(days)
        if row.get("event") == "page_view" and (space is None or row.get("space") == space)
    ]
    by_space = Counter(str(row.get("space", "unknown")) for row in events)
    by_day = Counter(str(row.get("date", "unknown")) for row in events)
    unique_visitors = {row["visitor_hash"] for row in events if row.get("visitor_hash")}
    return {
        "period_days": days,
        "space_filter": space,
        "page_views": len(events),
        "unique_visitors": len(unique_visitors),
        "by_space": dict(sorted(by_space.items())),
        "by_day": dict(sorted(by_day.items())),
    }


@api.get("/health")
def health() -> dict[str, Any]:
    return {
        "status": "ok",
        "storage_path": str(DATA_DIR),
        "storage_writable": os.access(DATA_DIR, os.W_OK),
    }


@api.post("/track", status_code=202)
def track(
    event: TrackEvent,
    request: Request,
    x_tracking_key: Annotated[str | None, Header()] = None,
) -> dict[str, str]:
    require_tracking_key(x_tracking_key)
    now = datetime.now(timezone.utc)
    client_ip = request.client.host if request.client else None
    append_event(
        {
            "timestamp": now.isoformat(),
            "date": now.date().isoformat(),
            "space": event.space,
            "path": event.path,
            "referrer": event.referrer,
            "visitor_hash": hash_visitor(event.visitor_id),
            "ip_hash": hash_client_ip(client_ip),
            "location": lookup_location(client_ip),
            "timezone": event.timezone,
            "language": event.language,
            "event": event.event,
            "user_agent": request.headers.get("user-agent", "")[:MAX_FIELD_LENGTH] or None,
        }
    )
    return {"status": "accepted"}


@api.get("/stats")
def stats(
    days: Annotated[int, Query(ge=1, le=366)] = 30,
    space: Annotated[str | None, Query(max_length=200)] = None,
    x_admin_key: Annotated[str | None, Header()] = None,
) -> dict[str, Any]:
    require_admin_key(x_admin_key)
    return calculate_stats(days, space)


@api.get("/events")
def recent_events(
    days: Annotated[int, Query(ge=1, le=30)] = 1,
    limit: Annotated[int, Query(ge=1, le=1000)] = 100,
    x_admin_key: Annotated[str | None, Header()] = None,
) -> dict[str, Any]:
    require_admin_key(x_admin_key)
    rows = list(iter_events(days))[-limit:]
    return {"count": len(rows), "events": rows}


TRACKER_JS = r"""
(() => {
  const script = document.currentScript;
  if (!script) return;
  const endpoint = script.dataset.endpoint;
  const space = script.dataset.space || window.location.hostname;
  const key = script.dataset.key || "";
  if (!endpoint) return;

  const storageKey = "hf_analytics_visitor";
  let visitorId = null;
  try {
    visitorId = localStorage.getItem(storageKey);
    if (!visitorId) {
      visitorId = globalThis.crypto?.randomUUID?.() ||
        `${Date.now()}-${Math.random().toString(36).slice(2)}`;
      localStorage.setItem(storageKey, visitorId);
    }
  } catch (_) {}

  const headers = {"Content-Type": "application/json"};
  if (key) headers["X-Tracking-Key"] = key;

  fetch(endpoint, {
    method: "POST",
    headers,
    body: JSON.stringify({
      space,
      path: window.location.pathname,
      referrer: document.referrer || null,
      visitor_id: visitorId,
    timezone: Intl.DateTimeFormat().resolvedOptions().timeZone || null,
    language: navigator.language || null,
      event: "page_view"
    }),
    keepalive: true,
    mode: "cors",
    credentials: "omit"
  }).catch(() => {});
})();
"""


@api.get("/tracker.js")
def tracker_js() -> Response:
    return Response(
        TRACKER_JS,
        media_type="application/javascript",
        headers={"Cache-Control": "public, max-age=300"},
    )


#@spaces.GPU(duration=1)
def dashboard_stats(days: int, space: str, admin_key: str):
    #if not torch.cuda.is_available():
        #raise RuntimeError("CUDA GPU is unavailable")

        #x = torch.randn(256, 256, device="cuda")
        #result = x @ x
        #torch.cuda.synchronize()

        #print(
            #f"GPU computation completed: "
            #f"{torch.cuda.get_device_name(0)}, device={result.device}",
            #flush=True,
        #)

    try:
        require_admin_key(admin_key or None)
        days = max(1, min(int(days), 366))
        result = calculate_stats(days, space.strip() or None)
        by_space = [[name, count] for name, count in sorted(
            result["by_space"].items(), key=lambda item: item[1], reverse=True
        )]
        by_day = [[day, count] for day, count in result["by_day"].items()]
        status = f"Showing the last {days} day(s)."
        return result["page_views"], result["unique_visitors"], by_space, by_day, status
    except HTTPException:
        raise gr.Error("Invalid admin key")
    except Exception as exc:
        raise gr.Error(f"Could not load statistics: {exc}") from exc


with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(
        "# 📊 HF Space Analytics\n"
        "Minimal page-view analytics stored in an attached Hugging Face Storage Bucket."
    )
    with gr.Row():
        days_input = gr.Number(value=30, minimum=1, maximum=366, precision=0, label="Days")
        space_input = gr.Textbox(label="Space filter", placeholder="Optional: username/space")
        admin_input = gr.Textbox(label="Admin key", type="password", placeholder="Only if configured")
    load_button = gr.Button("Load statistics", variant="primary")
    status_output = gr.Markdown()
    with gr.Row():
        views_output = gr.Number(label="Page views", interactive=False)
        visitors_output = gr.Number(label="Unique visitors", interactive=False)
    gr.Markdown("## Views by Space")
    spaces_output = gr.Dataframe(headers=["Space", "Views"], datatype=["str", "number"], interactive=False)
    gr.Markdown("## Views by Day")
    days_output = gr.Dataframe(headers=["Date", "Views"], datatype=["str", "number"], interactive=False)

    load_button.click(
        dashboard_stats,
        inputs=[days_input, space_input, admin_input],
        outputs=[views_output, visitors_output, spaces_output, days_output, status_output],
        api_name=False,
    )

# Custom API routes are registered first; Gradio is mounted last at the root.
app = gr.mount_gradio_app(api, demo, path="/")