import os
import json
import uuid
import time
import re
from pathlib import Path
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from gradio import Server
from openai import OpenAI

# Default system prompt for DeepSeek V4 Flash
DEFAULT_SYSTEM_PROMPT = "You are DeepSeek, a helpful AI assistant powered by DeepSeek V4 Flash. You provide accurate, detailed, and thoughtful responses."

# ── In-memory per-user session storage ──────────────────────────────────
USER_DATA = {}

def get_user_sessions(username):
    if username not in USER_DATA:
        USER_DATA[username] = {"sessions": {}, "active_session": None}
    return USER_DATA[username]

# ── Initialize Gradio Server (FastAPI subclass) ─────────────────────────
app = Server()

# Create static directory if it doesn't exist
STATIC_DIR = Path(__file__).parent / "static"
os.makedirs(STATIC_DIR, exist_ok=True)


# ── API Endpoint: Chat ──────────────────────────────────────────────────
@app.api(name="chat_with_deepseek")
def chat_with_deepseek(
    messages_json: str,
    reasoning_effort: str = "medium",
    max_tokens: str = "2048",
    temperature: str = "0.7",
    system_prompt: str = "",
) -> str:
    try:
        messages = json.loads(messages_json)
        max_tokens = int(max_tokens)
        temperature = float(temperature)

        key = os.environ.get("SILICONFLOW_API_KEY", "").strip()
        if not key:
            return json.dumps({
                "status": "error",
                "message": "SILICONFLOW_API_KEY environment variable is not configured on the server."
            })

        client = OpenAI(
            api_key=key,
            base_url="https://api.siliconflow.com/v1",
        )

        sys_prompt = system_prompt.strip() if system_prompt and system_prompt.strip() else DEFAULT_SYSTEM_PROMPT
        messages_no_system = [m for m in messages if m.get("role") != "system"]
        final_messages = [{"role": "system", "content": sys_prompt}] + messages_no_system

        params = {
            "model": "deepseek-ai/DeepSeek-V4-Flash",
            "messages": final_messages,
            "max_tokens": max_tokens,
            "temperature": temperature
        }

        if reasoning_effort in ["low", "medium", "high"]:
            params["reasoning_effort"] = reasoning_effort

        response = client.chat.completions.create(**params)

        content = response.choices[0].message.content
        reasoning_content = getattr(response.choices[0].message, "reasoning_content", "")

        if not reasoning_content and content and "<think" in content and "</think" in content:
            think_match = re.search(r'<think[^>]*>(.*?)</think', content, re.DOTALL)
            if think_match:
                reasoning_content = think_match.group(1).strip()
                content = re.sub(r'<think[^>]*>.*?</think\s*>', '', content, flags=re.DOTALL).strip()

        return json.dumps({
            "status": "success",
            "content": content,
            "reasoning_content": reasoning_content or ""
        })

    except Exception as e:
        return json.dumps({
            "status": "error",
            "message": str(e)
        })


# ── API Endpoint: Get User Info (OAuth) ────────────────────────────────
@app.api(name="get_user_info")
def get_user_info(profile) -> str:
    """Return the logged-in user's name, or empty string if not logged in."""
    # The Gradio Server passes the OAuth profile automatically
    if profile is None:
        return json.dumps({"logged_in": False, "username": ""})
    return json.dumps({"logged_in": True, "username": getattr(profile, 'name', '')})


# ── API Endpoint: Save Chat Session ────────────────────────────────────
@app.api(name="save_chat_session")
def save_chat_session(
    messages_json: str,
    title: str,
    session_id: str,
    profile = None,
) -> str:
    if profile is None:
        return json.dumps({"status": "error", "message": "Not logged in"})

    username = getattr(profile, 'name', '')
    ud = get_user_sessions(username)

    if not session_id or session_id == "new":
        sid = str(uuid.uuid4())[:8]
    else:
        sid = session_id

    messages = json.loads(messages_json) if messages_json else []

    if not title or title == "New Chat":
        for m in messages:
            if m.get("role") == "user":
                t = m.get("content", "")
                if isinstance(t, str):
                    title = t[:60] + ("..." if len(t) > 60 else "")
                break

    ud["sessions"][sid] = {
        "title": title,
        "messages": messages,
        "updated_at": time.time()
    }
    ud["active_session"] = sid

    return json.dumps({"status": "success", "session_id": sid, "title": title})


# ── API Endpoint: Load Chat Session ────────────────────────────────────
@app.api(name="load_chat_session")
def load_chat_session(
    session_id: str,
    profile = None,
) -> str:
    if profile is None:
        return json.dumps({"status": "error", "message": "Not logged in"})

    username = getattr(profile, 'name', '')
    ud = get_user_sessions(username)

    if session_id not in ud["sessions"]:
        return json.dumps({"status": "error", "message": "Session not found"})

    ud["active_session"] = session_id
    sess = ud["sessions"][session_id]

    return json.dumps({
        "status": "success",
        "session_id": session_id,
        "title": sess["title"],
        "messages": sess["messages"]
    })


# ── API Endpoint: List Chat Sessions ───────────────────────────────────
@app.api(name="list_chat_sessions")
def list_chat_sessions(profile = None) -> str:
    if profile is None:
        return json.dumps({"status": "error", "message": "Not logged in", "sessions": []})

    username = getattr(profile, 'name', '')
    ud = get_user_sessions(username)

    sessions = []
    for sid, sess in sorted(
        ud["sessions"].items(),
        key=lambda x: x[1].get("updated_at", 0),
        reverse=True
    ):
        sessions.append({
            "id": sid,
            "title": sess["title"],
            "updated_at": sess.get("updated_at", 0),
            "message_count": len(sess.get("messages", []))
        })

    return json.dumps({"status": "success", "sessions": sessions})


# ── API Endpoint: Delete Chat Session ──────────────────────────────────
@app.api(name="delete_chat_session")
def delete_chat_session(
    session_id: str,
    profile = None,
) -> str:
    if profile is None:
        return json.dumps({"status": "error", "message": "Not logged in"})

    username = getattr(profile, 'name', '')
    ud = get_user_sessions(username)

    if session_id in ud["sessions"]:
        del ud["sessions"][session_id]
        if ud["active_session"] == session_id:
            ud["active_session"] = None

    return json.dumps({"status": "success"})


# ── Serve the main HTML page ────────────────────────────────────────────
@app.get("/")
async def homepage():
    html_path = STATIC_DIR / "index.html"
    if html_path.exists():
        with open(html_path, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read(), status_code=200)
    return HTMLResponse(
        content="<h1>Frontend is building. Please refresh in a few seconds...</h1>",
        status_code=200
    )


# ── Mount static folder for CSS, JS, and image assets ──────────────────
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


# ── Launch ──────────────────────────────────────────────────────────────
if __name__ == "__main__":
    app.launch(show_error=True)