import os
import json
from typing import List, Tuple, Dict

import gradio as gr
from huggingface_hub import InferenceClient


# Default to GPT-OSS per your preference
HF_MODEL = os.getenv("HF_MODEL", "openai/gpt-oss-20b")


SYSTEM_PROMPT = (
    "You are Maya, the owner of Klinik Sehat Sentosa, a small outpatient clinic in Manado. "
    "A student systems analyst is interviewing you to gather information requirements for a simple "
    "appointment & queueing system (web + mobile).\n\n"
    "Goals: reduce patient wait time, prevent double bookings, support WhatsApp reminders, basic daily reports.\n"
    "Persona: friendly, busy, non-technical. Answer concretely based on realistic daily operations at a small clinic. "
    "If the student asks vague questions, ask specific clarifying questions before answering.\n"
    "Scope boundaries: No billing, no insurance, no EMR details—focus only on scheduling, queue order, reminders, and daily counts.\n"
    "Constraints: staff have low digital literacy; internet is intermittent; must run on existing Android phones; budget is small.\n"
    "Progress strategy: do not dump everything. Reveal details only when asked well. If the student asks leading questions, gently correct with realistic constraints.\n\n"
    "Style: Speak as Maya in first person. Be concise and concrete. Avoid technical jargon; describe operations in everyday terms."
)


def _build_messages(history: List[Tuple[str, str]], latest_user: str) -> List[Dict[str, str]]:
    messages: List[Dict[str, str]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    for user, assistant in history:
        if user:
            messages.append({"role": "user", "content": user})
        if assistant:
            messages.append({"role": "assistant", "content": assistant})
    if latest_user:
        messages.append({"role": "user", "content": latest_user})
    return messages


SUMMARY_SYSTEM_PROMPT = (
    "You are a requirements summarizer assisting the student and Maya. "
    "Based ONLY on the conversation transcript, produce a single JSON object that captures the current understanding of requirements. "
    "Do not invent details that were not stated or clearly implied by Maya. If unknown, use empty arrays or nulls.\n\n"
    "Output strictly valid JSON (no markdown, no extra text).\n\n"
    "Schema keys: \n"
    "- actors: string[]\n"
    "- goals: string[]\n"
    "- constraints: string[]\n"
    "- functional_requirements: string[]\n"
    "- non_functional_requirements: string[]\n"
    "- user_stories: string[] (format: 'As a <actor>, I want <need>, so that <benefit>.')\n"
    "- edge_cases: string[]\n"
    "- assumptions: string[]\n"
    "- open_questions: string[]\n"
    "- acceptance_criteria: string[]\n"
)


def _history_to_transcript(history: List[Tuple[str, str]]) -> str:
    lines = []
    for user, assistant in history:
        if user:
            lines.append(f"Student: {user}")
        if assistant:
            lines.append(f"Maya: {assistant}")
    return "\n".join(lines)


def user_submit(user_message: str, history: List[Tuple[str, str]]):
    history = history + [(user_message, "")]
    return "", history


def _oauth_or_env_token(hf_token: gr.OAuthToken | None):
    try:
        if hf_token and getattr(hf_token, "token", None):
            return hf_token.token
    except Exception:
        pass
    return os.getenv("HF_TOKEN")


def _provider_for_model(model_id: str) -> str | None:
    # Route provider explicitly when needed
    if model_id.startswith("openai/"):
        return "together"
    return None


def _provider_api_key(model_id: str) -> str | None:
    prov = _provider_for_model(model_id)
    if prov == "together":
        return os.getenv("TOGETHER_API_KEY")
    return None


def _inference_params(hf_token: gr.OAuthToken | None, model_id: str):
    provider = _provider_for_model(model_id)
    api_key = _provider_api_key(model_id)
    # Important: when using external providers (e.g., Together), do NOT pass HF OAuth/token,
    # otherwise the router attempts a delegated call and may 403.
    if provider:
        token = None
    else:
        token = _oauth_or_env_token(hf_token)
    return token, provider, api_key


def bot_reply(history: List[Tuple[str, str]], temperature: float = 0.7, top_p: float = 0.95, max_tokens: int = 512, hf_token: gr.OAuthToken = None):
    last_user = history[-1][0] if history else ""
    messages = _build_messages(history[:-1], last_user)
    token, provider, api_key = _inference_params(hf_token, HF_MODEL)
    client = InferenceClient(token=token, model=HF_MODEL)

    acc = ""
    try:
        for event in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
            provider=provider,
            api_key=api_key,
        ):
            token = ""
            try:
                choices = event.choices
                if len(choices) and getattr(choices[0], "delta", None) and choices[0].delta.content:
                    token = choices[0].delta.content
            except Exception:
                pass
            acc += token
            history[-1] = (history[-1][0], acc)
            yield history
    except Exception as e:
        err = (
            f"[Model error] {type(e).__name__}: {e}.\n"
            "Tip: Click Sign in or set HF_TOKEN.\n"
            "You can also set HF_MODEL to an HF-hosted chat model, e.g. 'meta-llama/Meta-Llama-3-8B-Instruct' or 'HuggingFaceH4/zephyr-7b-beta'."
        )
        history[-1] = (history[-1][0], err)
        yield history


def summarize(history: List[Tuple[str, str]], hf_token: gr.OAuthToken = None):
    transcript = _history_to_transcript(history)
    if not transcript.strip():
        return {"note": "No conversation yet. Ask Maya some questions first."}

    token, provider, api_key = _inference_params(hf_token, HF_MODEL)
    client = InferenceClient(token=token, model=HF_MODEL)
    messages = [
        {"role": "system", "content": SUMMARY_SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                "Using the transcript below, generate the JSON. Remember: valid JSON only.\n\n"
                f"Transcript:\n{transcript}"
            ),
        },
    ]
    # Non-streaming summarize for simplicity
    try:
        event = client.chat_completion(
            messages,
            max_tokens=1024,
            stream=False,
            temperature=0.2,
            top_p=0.95,
            provider=provider,
            api_key=api_key,
        )
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}"}
    # Extract content
    text = "{}"
    try:
        if event.choices and event.choices[0].message and event.choices[0].message.content:
            text = event.choices[0].message.content
    except Exception:
        pass
    try:
        return json.loads(text)
    except Exception:
        return {"raw": text}


with gr.Blocks(title="Maya – Klinik Owner (Role-Play)") as demo:
    # Top login section: enable OAuth on Spaces, skip locally
    is_space = bool(os.getenv("SPACE_ID") or os.getenv("SYSTEM") == "spaces")
    if is_space:
        gr.Markdown("## Sign in to use the model")
        gr.LoginButton()
    else:
        gr.Markdown("Tip: set env var `HF_TOKEN` for local use.")
    gr.Markdown("Model: `" + HF_MODEL + "`")

    gr.Markdown(
        """
        # Maya – Klinik Sehat Sentosa (Role-Play)
        Interview Maya to elicit requirements for an appointment & queueing system.
        - Ask clear, specific questions. Maya will ask for clarification if vague.
        - Click "Generate Requirements JSON" anytime to summarize current findings.
        """
    )

    with gr.Row():
        chatbot = gr.Chatbot(label="Maya (Clinic Owner)", height=420)
        json_out = gr.JSON(label="Requirements JSON", value=None)

    msg = gr.Textbox(placeholder="Ask Maya your next question…", label="Your message")

    with gr.Row():
        send_btn = gr.Button("Send", variant="primary")
        summarize_btn = gr.Button("Generate Requirements JSON")
        clear_btn = gr.ClearButton([chatbot, msg, json_out])

    # Send (Enter)
    msg.submit(user_submit, [msg, chatbot], [msg, chatbot]).then(
        bot_reply, [chatbot], [chatbot]
    )
    # Send (button)
    send_btn.click(user_submit, [msg, chatbot], [msg, chatbot]).then(
        bot_reply, [chatbot], [chatbot]
    )

    summarize_btn.click(summarize, [chatbot], [json_out])


if __name__ == "__main__":
    demo.launch()