import os, re, json
from functools import lru_cache

import gradio as gr
import torch

# ---------- Env/cache (quiet deprecation) ----------
os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
os.environ.setdefault("HF_HUB_CACHE", "/data/.cache/huggingface/hub")
os.environ.setdefault("GRADIO_TEMP_DIR", "/data/gradio")
os.environ.setdefault("GRADIO_CACHE_DIR", "/data/gradio")
os.environ.pop("TRANSFORMERS_CACHE", None)  # silence v5 deprecation note
for p in ["/data/.cache/huggingface/hub", "/data/gradio"]:
    try:
        os.makedirs(p, exist_ok=True)
    except Exception:
        pass

# ---------- Optional timezone ----------
try:
    from zoneinfo import ZoneInfo  # noqa: F401
except Exception:
    ZoneInfo = None  # noqa: N816

# ---------- Optional Cohere ----------
try:
    import cohere
    _HAS_COHERE = True
except Exception:
    _HAS_COHERE = False

from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login

# ---------- ClarityOps modules ----------
from safety import safety_filter, refusal_reply
from retriever import init_retriever, retrieve_context
from decision_math import compute_operational_numbers
from prompt_templates import build_system_preamble
from upload_ingest import extract_text_from_files
from session_rag import SessionRAG
from mdsi_analysis import capacity_projection, cost_estimate, outcomes_summary

# ---------- Config ----------
MODEL_ID = os.getenv("MODEL_ID", "CohereLabs/c4ai-command-r7b-12-2024")
HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") or os.getenv("HF_TOKEN")
COHERE_API_KEY = os.getenv("COHERE_API_KEY")
USE_HOSTED_COHERE = bool(COHERE_API_KEY and _HAS_COHERE)

# ---------- Helpers ----------
def pick_dtype_and_map():
    if torch.cuda.is_available():
        return torch.float16, "auto"
    if torch.backends.mps.is_available():
        return torch.float16, {"": "mps"}
    return torch.float32, "cpu"

def is_identity_query(message, history):
    patterns = [
        r"\bwho\s+are\s+you\b",
        r"\bwhat\s+are\s+you\b",
        r"\bwhat\s+is\s+your\s+name\b",
        r"\bwho\s+is\s+this\b",
        r"\bidentify\s+yourself\b",
        r"\btell\s+me\s+about\s+yourself\b",
        r"\bdescribe\s+yourself\b",
        r"\band\s+you\s*\?\b",
        r"\byour\s+name\b",
        r"\bwho\s+am\s+i\s+chatting\s+with\b",
    ]
    def match(t): return any(re.search(p, (t or "").strip().lower()) for p in patterns)
    if match(message):
        return True
    if history:
        last_user = history[-1][0] if isinstance(history[-1], (list, tuple)) else None
        if match(last_user):
            return True
    return False

def _iter_user_assistant(history):
    # history is a list of (user, assistant) tuples (Chatbot default format)
    for item in (history or []):
        if isinstance(item, (list, tuple)):
            u = item[0] if len(item) > 0 else ""
            a = item[1] if len(item) > 1 else ""
            yield u, a

def _history_to_prompt(message, history):
    parts = []
    for u, a in _iter_user_assistant(history):
        if u: parts.append(f"User: {u}")
        if a: parts.append(f"Assistant: {a}")
    parts.append(f"User: {message}")
    parts.append("Assistant:")
    return "\n".join(parts)

# ---------- Cohere path ----------
_co_client = None
if USE_HOSTED_COHERE:
    _co_client = cohere.Client(api_key=COHERE_API_KEY)

def cohere_chat(message, history):
    try:
        prompt = _history_to_prompt(message, history)
        resp = _co_client.chat(
            model="command-r7b-12-2024",
            message=prompt,
            temperature=0.3,
            max_tokens=900,
        )
        if hasattr(resp, "text") and resp.text: return resp.text.strip()
        if hasattr(resp, "reply") and resp.reply: return resp.reply.strip()
        if hasattr(resp, "generations") and resp.generations: return resp.generations[0].text.strip()
        return "Sorry, I couldn't parse the response from Cohere."
    except Exception as e:
        return f"Error calling Cohere API: {e}"

# ---------- Local model ----------
@lru_cache(maxsize=1)
def load_local_model():
    if not HF_TOKEN:
        raise RuntimeError("HUGGINGFACE_HUB_TOKEN is not set.")
    login(token=HF_TOKEN, add_to_git_credential=False)
    dtype, device_map = pick_dtype_and_map()
    tok = AutoTokenizer.from_pretrained(
        MODEL_ID, token=HF_TOKEN, use_fast=True, model_max_length=8192, padding_side="left", trust_remote_code=True,
    )
    mdl = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, token=HF_TOKEN, device_map=device_map, low_cpu_mem_usage=True,
        torch_dtype=dtype, trust_remote_code=True,
    )
    if mdl.config.eos_token_id is None and tok.eos_token_id is not None:
        mdl.config.eos_token_id = tok.eos_token_id
    return mdl, tok

def build_inputs(tokenizer, message, history):
    # Convert tuple history to chat template input for HF models
    msgs = []
    for u, a in _iter_user_assistant(history):
        if u: msgs.append({"role": "user", "content": u})
        if a: msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": message})
    return tokenizer.apply_chat_template(
        msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt"
    )

def local_generate(model, tokenizer, input_ids, max_new_tokens=900):
    input_ids = input_ids.to(model.device)
    with torch.no_grad():
        out = model.generate(
            input_ids=input_ids, max_new_tokens=max_new_tokens,
            do_sample=True, temperature=0.3, top_p=0.9,
            repetition_penalty=1.15,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    gen_only = out[0, input_ids.shape[-1]:]
    return tokenizer.decode(gen_only, skip_special_tokens=True).strip()

# ---------- Snapshot loader ----------
def _load_snapshot(path="snapshots/current.json"):
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        # Safe fallback if no snapshot present
        return {
            "timestamp": None, "beds_total": 400, "staffed_ratio": 1.0, "occupied_pct": 0.97,
            "ed_census": 62, "ed_admits_waiting": 19, "avg_ed_wait_hours": 8,
            "discharge_ready_today": 11, "discharge_barriers": {"allied_health": 7, "placement": 4},
            "rn_shortfall": {"med_ward_A": 1, "med_ward_B": 1},
            "forecast_admits_next_24h": {"respiratory": 14, "other": 9},
            "isolation_needs_waiting": {"contact": 3, "airborne": 1}, "telemetry_needed_waiting": 5
        }

# ---------- Init retrieval engines ----------
init_retriever()
_session_rag = SessionRAG()  # ephemeral per-session index for uploaded docs/images

# ---------- Executive pre-compute (MDSi block) ----------
def _mdsi_block():
    base_capacity = capacity_projection(18, 48, 6)
    cons_capacity = capacity_projection(12, 48, 6)
    opt_capacity = capacity_projection(24, 48, 6)
    cost_1200 = cost_estimate(1200, 74.0, 75000.0)
    outcomes = outcomes_summary()
    return json.dumps({
        "capacity_projection": {
            "conservative": cons_capacity, "base": base_capacity, "optimistic": opt_capacity
        },
        "cost_for_1200": cost_1200,
        "outcomes_summary": outcomes
    }, indent=2)

# ---------- Core chat logic ----------
def clarityops_reply(user_msg, history, tz, uploaded_files_paths):
    """
    - user_msg: latest message text
    - history: list[(user, assistant)]
    - tz: timezone str (unused but kept for future features)
    - uploaded_files_paths: list[str] absolute paths of uploaded files
    """
    try:
        # Safety (input)
        safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
        if blocked_in:
            return history + [(user_msg, refusal_reply(reason_in))]

        # Identity short-circuit
        if is_identity_query(safe_in, history):
            return history + [(user_msg, "I am ClarityOps, your strategic decision making AI partner.")]

        # Ingest new uploads into session RAG (ephemeral for this chat)
        if uploaded_files_paths:
            items = extract_text_from_files(uploaded_files_paths)
            if items:
                _session_rag.add_docs(items)

        # Pull session snippets from uploaded docs/images
        session_snips = "\n---\n".join(_session_rag.retrieve(
            "diabetes screening Indigenous Métis mobile program cost throughput outcomes logistics bed flow staffing discharge forecast",
            k=6
        ))

        # Load daily snapshot + policies + computed ops numbers
        snapshot = _load_snapshot()
        policy_context = retrieve_context(
            "mobile diabetes screening Indigenous community outreach logistics referral pathways cultural safety data governance cost effectiveness outcomes bed management discharge acceleration ambulance offload"
        )
        computed = compute_operational_numbers(snapshot)

        # Smart scenario detection: if user message suggests exec MDSi context, include pre-compute block
        user_lower = (safe_in or "").lower()
        mdsi_extra = _mdsi_block() if ("diabetes" in user_lower or "mdsi" in user_lower or "mobile screening" in user_lower) else ""

        system_preamble = build_system_preamble(
            snapshot=snapshot,
            policy_context=policy_context,
            computed_numbers=computed,
            scenario_text=(safe_in if len(safe_in) > 400 else "") + (f"\n\nExecutive Pre-Computed Blocks:\n{mdsi_extra}" if mdsi_extra else ""),
            session_snips=session_snips
        )

        augmented_user = system_preamble + "\n\nUser question or request:\n" + safe_in

        # Generate
        if USE_HOSTED_COHERE:
            out = cohere_chat(augmented_user, history)
        else:
            model, tokenizer = load_local_model()
            inputs = build_inputs(tokenizer, augmented_user, history)
            out = local_generate(model, tokenizer, inputs, max_new_tokens=900)

        # Tidy echoes
        if isinstance(out, str):
            for tag in ("Assistant:", "System:", "User:"):
                if out.startswith(tag):
                    out = out[len(tag):].strip()

        # Safety (output)
        safe_out, blocked_out, reason_out = safety_filter(out, mode="output")
        if blocked_out:
            out = refusal_reply(reason_out)

        return history + [(user_msg, safe_out)]
    except Exception as e:
        return history + [(user_msg, f"Error: {e}")]

# ---------- Theme & CSS ----------
theme = gr.themes.Soft(primary_hue="teal", neutral_hue="slate", radius_size=gr.themes.sizes.radius_lg)
custom_css = """
:root { --brand-bg: #e6f7f8; --brand-accent: #0d9488; --brand-text: #0f172a; --brand-text-light: #ffffff; }
.gradio-container { background: var(--brand-bg); }

/* Title */
h1 { color: var(--brand-text); font-weight: 700; font-size: 28px !important; }

/* Hide default Chatbot label */
.chatbot header, .chatbot .label, .chatbot .label-wrap, .chatbot .top, .chatbot .header, .chatbot > .wrap > header {
  display: none !important;
}

/* Chat bubbles */
.message.user, .message.bot {
  background: var(--brand-accent) !important;
  color: var(--brand-text-light) !important;
  border-radius: 12px !important;
  padding: 8px 12px !important;
}

/* Inputs softer */
textarea, input, .gr-input { border-radius: 12px !important; }
"""

# ---------- UI (single integrated window; uploads at bottom) ----------
with gr.Blocks(theme=theme, css=custom_css) as demo:
    # timezone capture (hidden)
    tz_box = gr.Textbox(visible=False)
    demo.load(
        lambda tz: tz,
        inputs=[tz_box],
        outputs=[tz_box],
        js="() => Intl.DateTimeFormat().resolvedOptions().timeZone",
    )

    # extra DOM cleanup for some gradio builds
    hide_label_sink = gr.HTML(visible=False)
    demo.load(
        fn=lambda: "",
        inputs=None,
        outputs=hide_label_sink,
        js="""
        () => {
          const sel = [
            '.chatbot header','.chatbot .label','.chatbot .label-wrap',
            '.chatbot .top','.chatbot .header','.chatbot > .wrap > header'
          ];
          sel.forEach(s => document.querySelectorAll(s).forEach(el => el.style.display = 'none'));
          return "";
        }
        """,
    )

    gr.Markdown("# ClarityOps Augmented Decision AI")

    # Main chat area (IMPORTANT: no type="messages" -> uses tuple history)
    chat = gr.Chatbot(label="", show_label=False, height=700)

    # ---- Bottom bar: uploads + message box + send/clear ----
    with gr.Row():
        uploads = gr.Files(
            label="Upload docs/images (PDF, DOCX, CSV, PNG, JPG)",
            file_types=["file"],
            file_count="multiple",
            height=68
        )

    with gr.Row():
        msg = gr.Textbox(
            label="",
            show_label=False,
            placeholder="Type a message… (paste scenarios here too; ClarityOps will adapt)",
            scale=10
        )
        send = gr.Button("Send", scale=1)
        clear = gr.Button("Clear chat", scale=1)

    # States
    state_history = gr.State(value=[])
    state_uploaded = gr.State(value=[])

    # When user selects files, store their paths in state (so they persist across turns)
    def _store_uploads(files, current):
        paths = []
        for f in (files or []):
            paths.append(getattr(f, "name", None) or f)
        return (current or []) + paths

    uploads.change(fn=_store_uploads, inputs=[uploads, state_uploaded], outputs=state_uploaded)

    # Send message -> compute reply -> update chat & history
    def _on_send(user_msg, history, tz, up_paths):
        if not user_msg or not user_msg.strip():
            return history, "", history  # no-op
        new_history = clarityops_reply(user_msg.strip(), history or [], tz, up_paths or [])
        return new_history, "", new_history

    send.click(
        fn=_on_send,
        inputs=[msg, state_history, tz_box, state_uploaded],
        outputs=[chat, msg, state_history],
        queue=True,
    )

    # Also allow pressing Enter inside the textbox
    msg.submit(
        fn=_on_send,
        inputs=[msg, state_history, tz_box, state_uploaded],
        outputs=[chat, msg, state_history],
        queue=True,
    )

    # Clear chat (keeps uploads so you can keep referencing docs)
    def _clear_chat():
        return [], [], []
    # Clear only chat + input; keep uploads
    clear.click(lambda: ([], "", []), None, [chat, msg, state_history])

if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    demo.launch(server_name="0.0.0.0", server_port=port, show_api=False, max_threads=8)