import gradio as gr
import json
import shutil
import tempfile
from pathlib import Path
import knowledge
from knowledge import (
    agent,
    agent_ctx,
    memory,
    run_with_trace,
    load_memory,
    reload_memory_sync,
    MEMORY_FILE,
)

PIN = "1969"


def get_text(content):
    """Extract plain text from Gradio 6 content blocks."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text")
    return str(content)


def check_pin(pin_input):
    """Verify PIN."""
    if not pin_input:
        return gr.update(visible=False), gr.update(visible=True), "Enter PIN"
    if pin_input == PIN:
        return gr.update(visible=True), gr.update(visible=False), ""
    return gr.update(visible=False), gr.update(visible=True), "Wrong PIN"


def format_history():
    """Load and format conversation history for display."""
    try:
        mem = load_memory()
        sessions = mem.get("sessions", [])
        if not sessions:
            return "No conversation history yet."

        output = []
        for i, ex in enumerate(reversed(sessions[-50:])):
            output.append(f"[{len(sessions)-i}] {ex.get('timestamp', '?')}")
            user_text = ex.get('user', '')[:200]
            agent_text = ex.get('agent', '')[:500]
            output.append(f"You: {user_text}")
            output.append(f"Brain: {agent_text}")
            if ex.get('searches'):
                output.append(f"Searches: {', '.join(ex['searches'])}")
            output.append("-" * 40)
        return "\n".join(output)
    except Exception as e:
        return f"Error loading history: {str(e)}"


def export_memory():
    """Export memory.json for download."""
    if not MEMORY_FILE.exists():
        return None
    temp_dir = tempfile.mkdtemp()
    temp_path = Path(temp_dir) / "brain_memory.json"
    shutil.copy(MEMORY_FILE, temp_path)
    return temp_path


def import_memory(file):
    """Import memory.json from upload."""
    if file is None:
        return "No file uploaded"

    try:
        uploaded_path = Path(file.name)
        with open(uploaded_path, 'r', encoding='utf-8') as f:
            data = json.load(f)

        if "sessions" not in data:
            return "Invalid memory file: missing 'sessions'"

        if MEMORY_FILE.exists():
            backup_path = MEMORY_FILE.with_suffix('.backup.json')
            shutil.copy(MEMORY_FILE, backup_path)

        shutil.copy(uploaded_path, MEMORY_FILE)

        # Reload memory
        global memory
        memory = reload_memory_sync()
        knowledge.memory = memory

        session_count = len(memory.get("sessions", []))
        return f"Memory imported successfully! {session_count} sessions loaded."
    except Exception as e:
        return f"Error importing: {str(e)}"


async def chat(message, history):
    """Main chat handler."""
    user_text = get_text(message)

    if not user_text or not user_text.strip():
        return ""

    if agent is None or agent_ctx is None:
        return "Error: System failed to initialize. Check that DEEPSEEK_API_KEY is set."

    try:
        reasoning, response, searches = await run_with_trace(
            agent, agent_ctx, user_text, knowledge.memory
        )

        # Build output with reasoning trace
        output = ""

        if searches:
            output += "**Searches:**\n"
            for s in searches:
                output += f"- `{s}`\n"
            output += "\n"

        if reasoning:
            output += "**Reasoning:**\n"
            for step in reasoning:
                if step["type"] == "tool_call":
                    output += f"> **Tool:** {step['tool']}\n"
                    output += f"> **Found:** {step['result_preview']}\n\n"
            output += "---\n\n"

        output += response
        return output

    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"Error: {str(e)}"


# Build UI
with gr.Blocks(title="Brain") as demo:
    gr.Markdown("# Brain")

    with gr.Tabs():
        with gr.TabItem("Chat"):
            chatbot = gr.ChatInterface(
                fn=chat,
                chatbot=gr.Chatbot(height=500),
            )

        with gr.TabItem("History"):
            gr.Markdown("View conversation history (PIN protected)")

            with gr.Column(visible=True) as pin_section:
                pin_input = gr.Textbox(
                    label="PIN",
                    type="password",
                    max_lines=1,
                )
                with gr.Row():
                    pin_btn = gr.Button("Unlock")
                    import_btn = gr.Button("Import Memory")
                pin_error = gr.Markdown("")
                import_file = gr.File(label="Upload memory.json", visible=False)
                import_status = gr.Markdown("")

            with gr.Column(visible=False) as history_section:
                with gr.Row():
                    refresh_btn = gr.Button("Refresh")
                    export_btn = gr.Button("Export Memory")
                    lock_btn = gr.Button("Lock")
                history_display = gr.Textbox(
                    label="Conversation History",
                    lines=30,
                    max_lines=50,
                    interactive=False,
                )
                export_file = gr.File(label="Download Memory")

            def lock_history():
                return gr.update(visible=False), gr.update(visible=True), ""

            pin_btn.click(
                fn=check_pin,
                inputs=[pin_input],
                outputs=[history_section, pin_section, pin_error]
            )
            refresh_btn.click(fn=format_history, outputs=[history_display])
            pin_btn.click(fn=format_history, outputs=[history_display])
            lock_btn.click(
                fn=lock_history,
                outputs=[history_section, pin_section, pin_error]
            )
            export_btn.click(fn=export_memory, outputs=[export_file])
            import_btn.click(lambda: gr.update(visible=True), outputs=[import_file])
            import_file.change(fn=import_memory, inputs=[import_file], outputs=[import_status])
            pin_btn.click(lambda: "", outputs=[pin_input])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=5)
    demo.launch(show_error=True)