import json
import os
import traceback
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import gradio as gr

# Adjust this import to match your package layout.
# If PackedLLMRunner lives in a local module, change this line accordingly.
from PackedLLM import PackedLLMRunner


CHECKPOINT_PATH = os.getenv("PACKEDLLM_CHECKPOINT", "PackedLLM.pt")
MAP_LOCATION = os.getenv("PACKEDLLM_MAP_LOCATION", "cpu")
BOT_ID = os.getenv("PACKEDLLM_BOT_ID", "pip")
USER_ID = os.getenv("PACKEDLLM_USER_ID", "space_user")
DEFAULT_VENV_ID = os.getenv("PACKEDLLM_VENV_ID", "space_default")


_runner: Optional[PackedLLMRunner] = None


def _format_any(value: Any) -> str:
    if isinstance(value, str):
        return value
    try:
        return json.dumps(value, indent=2, ensure_ascii=False, default=str)
    except Exception:
        return repr(value)


def _parse_json(text: str, fallback: Any = None) -> Any:
    text = (text or "").strip()
    if not text:
        return fallback
    return json.loads(text)


def _parse_tags(text: str) -> List[str]:
    raw = (text or "").strip()
    if not raw:
        return []
    return [t.strip() for t in raw.split(",") if t.strip()]


def get_runner() -> PackedLLMRunner:
    global _runner
    if _runner is None:
        _runner = PackedLLMRunner(
            checkpoint_path=CHECKPOINT_PATH,
            map_location=MAP_LOCATION,
            bot_id=BOT_ID,
            user_id=USER_ID,
            warmup=False,
            verbose=True,
        )
    return _runner


def chat_fn(
    message: str,
    history: List[Dict[str, str]],
    image: Optional[str],
    deep_think: bool,
    fast_think: bool,
    stream: bool,
):
    history = history or []
    message = (message or "").strip()

    if not message and not image:
        yield history, history, ""
        return

    user_content = message
    if image:
        if user_content:
            user_content += f"\n\n[Image attached: {image}]"
        else:
            user_content = f"[Image attached: {image}]"

    history = history + [{"role": "user", "content": user_content}]
    yield history, history, ""

    try:
        runner = get_runner()
        result = runner.chat(
            message,
            image=image,
            stream=stream,
            deep_think=deep_think,
            fast_think=fast_think,
        )

        if stream and hasattr(result, "__iter__") and not isinstance(result, (str, bytes, dict, list, tuple)):
            assembled = ""
            for chunk in result:
                if isinstance(chunk, dict) and "content" in chunk:
                    chunk_text = str(chunk["content"])
                else:
                    chunk_text = str(chunk)
                assembled += chunk_text
                live_history = history + [{"role": "assistant", "content": assembled}]
                yield live_history, live_history, ""

            history = history + [{"role": "assistant", "content": assembled or ""}]
            yield history, history, ""
        else:
            history = history + [{"role": "assistant", "content": _format_any(result)}]
            yield history, history, ""

    except Exception as exc:
        err = traceback.format_exc()
        history = history + [
            {
                "role": "assistant",
                "content": f"**Error:** {exc}\n\n```text\n{err}\n```",
            }
        ]
        yield history, history, ""


def run_expert(
    expert_name: str,
    prompt: str,
    image: Optional[str],
    character_card: str,
    logic_mode: str,
    tools_json: str,
) -> str:
    runner = get_runner()
    prompt = (prompt or "").strip()

    if not prompt and expert_name not in {"vision", "tool"}:
        raise gr.Error("Please enter a prompt.")

    kwargs: Dict[str, Any] = {}

    if expert_name == "role" and character_card.strip():
        kwargs["character_card"] = character_card.strip()

    if expert_name == "logic":
        kwargs["mode"] = logic_mode or "deep_then_answer"

    if expert_name == "vision":
        if not image:
            raise gr.Error("Please upload an image for the Vision expert.")
        kwargs["image"] = image

    if expert_name == "tool":
        tools = _parse_json(tools_json, fallback=None)
        if not tools:
            tools = [
                {
                    "name": "noop",
                    "description": "No-op demo tool.",
                    "parameters": {"type": "object", "properties": {}},
                }
            ]
        return _format_any(runner.tool(prompt, tools=tools))

    if expert_name == "head":
        return _format_any(runner.head(prompt, image=image, **kwargs))
    if expert_name == "creative":
        return _format_any(runner.creative(prompt, **kwargs))
    if expert_name == "code":
        return _format_any(runner.code(prompt, **kwargs))
    if expert_name == "logic":
        return _format_any(runner.logic(prompt, **kwargs))
    if expert_name == "math":
        return _format_any(runner.math(prompt, **kwargs))
    if expert_name == "translate":
        return _format_any(runner.translate(prompt, **kwargs))
    if expert_name == "affect":
        return _format_any(runner.affect(prompt, **kwargs))
    if expert_name == "role":
        return _format_any(runner.role(prompt, **kwargs))
    if expert_name == "vision":
        return _format_any(runner.vision(prompt, image=image, **kwargs))
    if expert_name == "web":
        return _format_any(runner.web(prompt, **kwargs))
    if expert_name == "action":
        return _format_any(runner.action(prompt, **kwargs))

    raise gr.Error(f"Unknown expert: {expert_name}")


def store_memory(text: str, tags: str, importance: float) -> str:
    runner = get_runner()
    result = runner.memory_store(
        text=text.strip(),
        tags=_parse_tags(tags),
        importance=float(importance),
    )
    return _format_any(result)


def recall_memory(query: str, top_k: int) -> str:
    runner = get_runner()
    result = runner.memory_recall(query.strip(), top_k=int(top_k))
    return _format_any(result)


def refresh_profiles() -> tuple[str, str]:
    runner = get_runner()
    return (
        json.dumps(runner.get_user_profile(), indent=2, ensure_ascii=False, default=str),
        json.dumps(runner.get_bot_profile(), indent=2, ensure_ascii=False, default=str),
    )


def apply_profiles(user_profile_json: str, bot_profile_json: str) -> str:
    runner = get_runner()
    user_updates = _parse_json(user_profile_json, fallback={})
    bot_updates = _parse_json(bot_profile_json, fallback={})

    if not isinstance(user_updates, dict):
        raise gr.Error("User profile JSON must be an object.")
    if not isinstance(bot_updates, dict):
        raise gr.Error("Bot profile JSON must be an object.")

    if user_updates:
        runner.set_user_profile(user_updates)
    if bot_updates:
        runner.set_bot_profile(bot_updates)

    return "Profiles updated."


def warmup_runner(include_web: bool, include_vision: bool, include_action: bool) -> str:
    runner = get_runner()
    report = runner.warmup(
        include_web=include_web,
        include_vision=include_vision,
        include_action=include_action,
    )
    return _format_any(report)


def get_status() -> str:
    runner = get_runner()
    return _format_any(runner.status())


def reload_expert(expert_name: str) -> str:
    runner = get_runner()
    result = runner.reload_expert(expert_name)
    return _format_any(result)


def unload_expert(expert_name: str) -> str:
    runner = get_runner()
    runner.unload_expert(expert_name)
    return f"Unloaded {expert_name}."


def unload_all() -> str:
    runner = get_runner()
    runner.unload_all()
    return "Unloaded all experts."


def save_checkpoint(path: str) -> str:
    runner = get_runner()
    path = (path or "").strip()
    if path:
        runner.save(path)
        return f"Saved checkpoint to {path}"
    runner.save()
    return f"Saved checkpoint to {CHECKPOINT_PATH}"


def run_code(
    code: str,
    venv_id: str,
    requirements_text: str,
    timeout: int,
    max_ram_mb: int,
) -> str:
    runner = get_runner()
    reqs = [line.strip() for line in (requirements_text or "").splitlines() if line.strip()]
    result = runner.run_code(
        code=code,
        venv_id=venv_id or DEFAULT_VENV_ID,
        requirements=reqs or None,
        timeout=int(timeout),
        max_ram_mb=int(max_ram_mb),
        ensure_venv=True,
    )
    return _format_any(result)


def web_search(query: str, deep_search: bool) -> str:
    runner = get_runner()
    result = runner.web_search(query.strip(), deep_search=deep_search)
    return _format_any(result)


with gr.Blocks(title="PackedLLM Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        f"""
# PackedLLM

Loaded checkpoint: `{CHECKPOINT_PATH}`

This demo exposes the main chat pipeline, direct expert calls, memory, web search, code execution, and system controls.
"""
    )

    with gr.Tabs():
        with gr.Tab("Chat"):
            chatbot = gr.Chatbot(type="messages", height=600, label="PackedLLM Chat")
            history_state = gr.State([])

            with gr.Row():
                image_in = gr.Image(
                    type="filepath",
                    label="Optional image for vision-enabled turns",
                )

            prompt_in = gr.Textbox(
                label="Message",
                placeholder="Ask PackedLLM anything...",
                lines=3,
            )

            with gr.Row():
                deep_think_in = gr.Checkbox(value=False, label="Deep think")
                fast_think_in = gr.Checkbox(value=False, label="Fast think")
                stream_in = gr.Checkbox(value=True, label="Stream")

            with gr.Row():
                send_btn = gr.Button("Send", variant="primary")
                clear_btn = gr.Button("Clear")

            send_btn.click(
                chat_fn,
                inputs=[prompt_in, history_state, image_in, deep_think_in, fast_think_in, stream_in],
                outputs=[chatbot, history_state, prompt_in],
            )
            prompt_in.submit(
                chat_fn,
                inputs=[prompt_in, history_state, image_in, deep_think_in, fast_think_in, stream_in],
                outputs=[chatbot, history_state, prompt_in],
            )

            def clear_chat():
                return [], [], ""

            clear_btn.click(clear_chat, outputs=[chatbot, history_state, prompt_in])

        with gr.Tab("Experts"):
            gr.Markdown("Call individual experts directly.")

            with gr.Row():
                expert = gr.Dropdown(
                    choices=[
                        "head",
                        "creative",
                        "code",
                        "logic",
                        "math",
                        "translate",
                        "affect",
                        "role",
                        "vision",
                        "tool",
                        "web",
                        "action",
                    ],
                    value="head",
                    label="Expert",
                )
                logic_mode = gr.Dropdown(
                    choices=["deep_then_answer", "answer_only", "deep_only"],
                    value="deep_then_answer",
                    label="Logic mode",
                )

            expert_prompt = gr.Textbox(label="Prompt", lines=6, placeholder="Enter a prompt for the chosen expert.")
            expert_image = gr.Image(type="filepath", label="Image for Vision expert (optional)")
            character_card = gr.Textbox(
                label="Character card for Role expert",
                lines=4,
                placeholder="You are Pip, a direct and slightly sarcastic assistant.",
            )
            tools_json = gr.Textbox(
                label="Tools JSON for Tool expert",
                lines=8,
                placeholder='[{"name":"noop","description":"No-op demo tool.","parameters":{"type":"object","properties":{}}}]',
            )
            expert_out = gr.Textbox(label="Output", lines=18)

            run_expert_btn = gr.Button("Run expert", variant="primary")
            run_expert_btn.click(
                run_expert,
                inputs=[expert, expert_prompt, expert_image, character_card, logic_mode, tools_json],
                outputs=[expert_out],
            )

        with gr.Tab("Memory"):
            gr.Markdown("Store and recall memory, plus user/bot profile editing.")

            with gr.Row():
                mem_text = gr.Textbox(label="Text to store", lines=4)
                mem_tags = gr.Textbox(label="Tags (comma-separated)", value="manual")
            mem_importance = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Importance")
            store_btn = gr.Button("Store memory")
            mem_store_out = gr.Textbox(label="Store result", lines=4)

            with gr.Row():
                mem_query = gr.Textbox(label="Recall query", lines=3)
                mem_top_k = gr.Slider(1, 20, value=5, step=1, label="Top K")
            recall_btn = gr.Button("Recall memory")
            mem_recall_out = gr.Textbox(label="Recall result", lines=10)

            gr.Markdown("Profiles")
            with gr.Row():
                user_profile_json = gr.Textbox(label="User profile JSON", lines=10)
                bot_profile_json = gr.Textbox(label="Bot profile JSON", lines=10)
            with gr.Row():
                refresh_profiles_btn = gr.Button("Refresh profiles")
                apply_profiles_btn = gr.Button("Apply profiles", variant="primary")
            profile_status = gr.Textbox(label="Profile status", lines=2)

            store_btn.click(
                store_memory,
                inputs=[mem_text, mem_tags, mem_importance],
                outputs=[mem_store_out],
            )
            recall_btn.click(
                recall_memory,
                inputs=[mem_query, mem_top_k],
                outputs=[mem_recall_out],
            )
            refresh_profiles_btn.click(
                refresh_profiles,
                inputs=[],
                outputs=[user_profile_json, bot_profile_json],
            )
            apply_profiles_btn.click(
                apply_profiles,
                inputs=[user_profile_json, bot_profile_json],
                outputs=[profile_status],
            )

        with gr.Tab("Web"):
            gr.Markdown("Direct web search through the embedded web module.")
            web_query = gr.Textbox(label="Search query", lines=3)
            web_deep = gr.Checkbox(value=False, label="Deep search")
            web_btn = gr.Button("Search", variant="primary")
            web_out = gr.Textbox(label="Results", lines=20)
            web_btn.click(web_search, inputs=[web_query, web_deep], outputs=[web_out])

        with gr.Tab("CodeBox"):
            gr.Markdown("Run code inside the embedded sandbox.")
            code_text = gr.Code(label="Python code", language="python", lines=18)
            with gr.Row():
                venv_id_in = gr.Textbox(label="Venv ID", value=DEFAULT_VENV_ID)
                timeout_in = gr.Slider(5, 600, value=120, step=5, label="Timeout (seconds)")
                max_ram_in = gr.Slider(256, 32768, value=4096, step=256, label="Max RAM (MB)")
            requirements_in = gr.Textbox(
                label="Requirements (one per line)",
                lines=5,
                placeholder="numpy\npandas\nrequests",
            )
            code_btn = gr.Button("Run code", variant="primary")
            code_out = gr.Textbox(label="Sandbox result", lines=20)
            code_btn.click(
                run_code,
                inputs=[code_text, venv_id_in, requirements_in, timeout_in, max_ram_in],
                outputs=[code_out],
            )

        with gr.Tab("System"):
            gr.Markdown("Warmup, status, load management, and checkpoint saving.")

            with gr.Row():
                warm_web = gr.Checkbox(value=False, label="Warm web")
                warm_vision = gr.Checkbox(value=False, label="Warm vision")
                warm_action = gr.Checkbox(value=False, label="Warm action")
            warm_btn = gr.Button("Warmup", variant="primary")
            warm_out = gr.Textbox(label="Warmup report", lines=10)

            status_btn = gr.Button("Refresh status")
            status_out = gr.Textbox(label="Status", lines=16)

            with gr.Row():
                expert_name = gr.Dropdown(
                    choices=[
                        "head_expert",
                        "creative_expert",
                        "code_expert",
                        "logic_expert",
                        "math_expert",
                        "affect_expert",
                        "role_expert",
                        "vision_expert",
                        "tool_expert",
                        "translation_expert",
                        "web_expert",
                        "action_expert",
                    ],
                    value="head_expert",
                    label="Expert to reload/unload",
                )
                save_path = gr.Textbox(label="Save path (blank = default)", value="")
            with gr.Row():
                reload_btn = gr.Button("Reload expert")
                unload_btn = gr.Button("Unload expert")
                unload_all_btn = gr.Button("Unload all")
                save_btn = gr.Button("Save checkpoint", variant="primary")

            reload_out = gr.Textbox(label="Reload result", lines=2)
            unload_out = gr.Textbox(label="Unload result", lines=2)
            save_out = gr.Textbox(label="Save result", lines=2)

            warm_btn.click(
                warmup_runner,
                inputs=[warm_web, warm_vision, warm_action],
                outputs=[warm_out],
            )
            status_btn.click(get_status, inputs=[], outputs=[status_out])
            reload_btn.click(reload_expert, inputs=[expert_name], outputs=[reload_out])
            unload_btn.click(unload_expert, inputs=[expert_name], outputs=[unload_out])
            unload_all_btn.click(unload_all, inputs=[], outputs=[unload_out])
            save_btn.click(save_checkpoint, inputs=[save_path], outputs=[save_out])

    demo.load(refresh_profiles, inputs=[], outputs=[user_profile_json, bot_profile_json])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()