import asyncio
import json
import subprocess
from pathlib import Path
from datetime import datetime, timezone, timedelta
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

def ensure_initial_message():
    msgs = load_messages()
    if len(msgs) == 0:
        greeting = {
            "role": "bot",
            "text": "Hi. I am Dwarakesh. Ask me anything about myself.",
            "timestamp": datetime.now(timezone(timedelta(hours=5, minutes=30))).replace(tzinfo=None).isoformat()
        }
        save_messages([greeting])

# --------------------------------------------------------
# FASTAPI INIT
# --------------------------------------------------------

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],      # allow all for HF Spaces
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --------------------------------------------------------
# JSON DATABASE (persistent message storage)
# --------------------------------------------------------

ROOT_DIR = Path(__file__).parent.parent
DATA_DIR = ROOT_DIR / "data"
DATA_DIR.mkdir(exist_ok=True)

DB_FILE = DATA_DIR / "history.json"


def load_messages():
    if not DB_FILE.exists():
        return []
    try:
        with open(DB_FILE, "r") as f:
            return json.load(f)
    except:
        return []


def save_messages(msgs):
    with open(DB_FILE, "w") as f:
        json.dump(msgs, f, indent=2)


def add_message(role, text, src=None):
    msgs = load_messages()
    entry = {
        "role": role,
        "text": text,
        "src": src,
        "timestamp": datetime.now(timezone(timedelta(hours=5, minutes=30))).replace(tzinfo=None).isoformat()
    }
    msgs.append(entry)
    save_messages(msgs)
    return entry


# --------------------------------------------------------
# PYTHON ENGINE SPAWN
# --------------------------------------------------------

python_engine_path = ROOT_DIR / "python_engine" / "engine_server.py"

python = subprocess.Popen(
    ["python3", str(python_engine_path)],  # Docker-safe, HF-safe
    stdin=subprocess.PIPE,
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE,
    text=True,
    bufsize=1
)

ready = False
callbacks = []


# --------------------------------------------------------
# ENGINE ASYNC READERS
# --------------------------------------------------------

async def read_stdout():
    """Reads JSON responses from python engine."""
    global callbacks
    loop = asyncio.get_event_loop()

    while True:
        line = await loop.run_in_executor(None, python.stdout.readline)
        if not line:
            continue

        line = line.strip()
        if callbacks:
            cb = callbacks.pop(0)
            try:
                cb(json.loads(line))
            except Exception as err:
                print("JSON decode error:", err)
                print("Line:", line)


async def read_stderr():
    """Reads stderr for logs and READY signal."""
    global ready
    loop = asyncio.get_event_loop()

    while True:
        line = await loop.run_in_executor(None, python.stderr.readline)
        if not line:
            continue

        msg = line.strip()
        print("[PY]", msg)

        if "READY" in msg:
            ready = True


@app.on_event("startup")
async def startup_event():
    ensure_initial_message()
    asyncio.create_task(read_stdout())
    asyncio.create_task(read_stderr())


# --------------------------------------------------------
# SEND QUERY TO PYTHON ENGINE
# --------------------------------------------------------

async def ask_python(query: str):
    if not ready:
        raise HTTPException(503, "Python engine not ready")

    loop = asyncio.get_event_loop()
    future = loop.create_future()

    callbacks.append(future.set_result)

    python.stdin.write(json.dumps({"query": query}) + "\n")
    python.stdin.flush()

    return await future


# --------------------------------------------------------
# API ROUTES
# --------------------------------------------------------

@app.post("/ask")
async def ask_route(payload: dict):
    user_query = payload.get("query")
    if not user_query:
        raise HTTPException(400, "Missing 'query' field")

    # Save user message
    add_message("user", user_query)

    # Engine response
    result = await ask_python(user_query)

    # Save bot responses
    if isinstance(result.get("response"), list):
        for msg in result["response"]:
            add_message("bot", msg)

    return result


@app.get("/history")
async def history_route():
    ensure_initial_message()
    return load_messages()


@app.post("/clear")
async def clear_route():
    msgs = []
    greeting = {
        "role": "bot",
        "text": "Hi. I am Dwarakesh. Ask me anything about myself.",
        "timestamp": datetime.now(timezone(timedelta(hours=5, minutes=30))).replace(tzinfo=None).isoformat()
    }
    msgs.append(greeting)
    save_messages(msgs)
    return {"success": True, "message": greeting}


# --------------------------------------------------------
# STATIC FRONTEND (Vite build)
# --------------------------------------------------------

FRONTEND_DIR = ROOT_DIR / "frontend"
DIST_DIR = FRONTEND_DIR / "dist"
PUBLIC_DIR = FRONTEND_DIR / "public"

# /static/* → public folder
app.mount("/static", StaticFiles(directory=PUBLIC_DIR), name="static")

# /assets/* → dist/assets
app.mount("/assets", StaticFiles(directory=DIST_DIR / "assets"), name="assets")

# favicon
@app.get("/favicon.png")
async def favicon():
    return FileResponse(DIST_DIR / "favicon.png")

# Serve frontend for all other routes
@app.get("/{full_path:path}")
async def serve_spa(full_path: str):
    return FileResponse(DIST_DIR / "index.html")