import os
import re
import gradio as gr
import spaces

try:
    from llama_cpp import Llama
    USE_LLAMA_CPP = True
    print("llama-cpp-python is installed. Will use GGUF inference.")
except ImportError:
    USE_LLAMA_CPP = False
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    print("llama-cpp-python is not installed. Falling back to transformers inference.")

if USE_LLAMA_CPP:
    GGUF_FILENAME = "production_mindmap_model.gguf"
    if os.path.exists(f"./{GGUF_FILENAME}"):
        GGUF_PATH = f"./{GGUF_FILENAME}"
    else:
        from huggingface_hub import hf_hub_download
        MODEL_REPO_ID = os.environ.get("MODEL_REPO_ID", "kazutab/mindmap-studio-model")
        if not MODEL_REPO_ID:
            raise ValueError("Local model not found and MODEL_REPO_ID is not set.")
        GGUF_PATH = hf_hub_download(repo_id=MODEL_REPO_ID, filename="backend/production_mindmap_model.gguf")
    
    model = Llama(model_path=GGUF_PATH, n_ctx=2048, n_gpu_layers=-1)
else:
    MERGED_MODEL_PATH = "./production_mindmap_model_merged"
    tokenizer = AutoTokenizer.from_pretrained(MERGED_MODEL_PATH)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(MERGED_MODEL_PATH, torch_dtype=torch.float16, device_map="auto")
    model.eval()

STRICT_SYSTEM_PROMPT = """あなたは極めて優秀で厳密な情報抽出アシスタントです。入力文章の論理構造を正確に読み取り、Markdown形式の目次（マインドマップ）を出力してください。

【厳格なルール】
1. 企業名、人物名、数値などの固有名詞や事実関係は、入力文章から「正確に」抽出すること。絶対に捏造したり、他の主語と混同したりしてはなりません。
2. ただし、情報を階層化して分かりやすく整理するための「一般的なカテゴリ名（例：背景、特徴、課題、概要など）」を親見出しとして補足することは許可します。
3. あなたの推測や外部知識を一切混ぜず、入力文章内の事実のみを構造化してください。"""

@spaces.GPU(duration=120)
def generate_mindmap(input_text: str):
    input_text = input_text.strip()
    if not input_text:
        return ""

    USER_PROMPT = f"""以下の文章から論理構造を抽出し、Markdown形式の目次（マインドマップ）を出力してください。
【出力時の厳守ルール（違反厳禁）】
1. 否定表現の厳守：「〜しない」「過度に依存しない」などの否定表現を絶対に見落とさず、意味を逆転させないこと。
2. 創作の禁止：記事に明記されていない具体的な行動や予定（例：「〜への参加」「〜の強化を目指す」など）を勝手に推測して付け足さないこと。
3. 事実の完全一致：抽出した内容が、元の文章の事実と完全に一致していることのみを出力すること。

入力文章:
{input_text}"""
    messages = [{"role": "system", "content": STRICT_SYSTEM_PROMPT}, {"role": "user", "content": USER_PROMPT}]
    
    if USE_LLAMA_CPP:
        response = model.create_chat_completion(
            messages=messages, max_tokens=1024, temperature=0.0, repeat_penalty=1.1
        )
        generated_markdown = response['choices'][0]['message']['content'].strip()
    else:
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=False, repetition_penalty=1.1)
        generated_markdown = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    
    generated_markdown = re.sub(r'\s*(#+ )', r'\n\1', generated_markdown).strip()
    if not generated_markdown.startswith('#'):
        if '##' in generated_markdown or '###' in generated_markdown:
            generated_markdown = "# マインドマップ\n" + generated_markdown
        else:
            generated_markdown = "# マインドマップ\n## 抽出結果\n- " + generated_markdown.replace('\n', '\n- ')

    return generated_markdown


# -----------------------------------------------------------------------------------------
# JS / CSS / HTML Injection for 100% Perfect Layout bypass
# -----------------------------------------------------------------------------------------

base_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600&display=swap');

:root, .dark, body {
    --bg-canvas: #fafafa !important;
    --bg-sidebar: #ffffff !important;
    --border-color: #eaeaea !important;
    --text-primary: #171717 !important;
    --text-secondary: #666666 !important;
    --text-tertiary: #a1a1aa !important;
    --focus-ring: rgba(0, 0, 0, 0.08) !important;
    --body-text-color: #171717 !important;
    color: #171717 !important;
}

* { box-sizing: border-box; margin: 0; padding: 0; }
body { font-family: -apple-system, BlinkMacSystemFont, "Inter", "Segoe UI", "Roboto", "Helvetica Neue", sans-serif; color: var(--text-primary); background-color: var(--bg-canvas); height: 100vh; overflow: hidden; -webkit-font-smoothing: antialiased; }

/* Force layout to break out of Gradio's container limitations */
.app-layout { position: fixed !important; top: 0 !important; left: 0 !important; display: flex !important; height: 100vh !important; width: 100vw !important; z-index: 99999 !important; background-color: var(--bg-canvas) !important; margin: 0 !important; padding: 0 !important; }

/* Force light mode text colors to defeat Gradio Dark Mode */
.app-layout, .app-layout div, .app-layout span, .app-layout p, .app-layout label { color: #171717 !important; }

.sidebar { width: 360px !important; min-width: 360px !important; max-width: 360px !important; background-color: var(--bg-sidebar) !important; border-right: 1px solid var(--border-color) !important; display: flex !important; flex-direction: column !important; box-shadow: 1px 0 10px rgba(0,0,0,0.02) !important; z-index: 10 !important; height: 100vh !important; margin: 0 !important; padding: 0 !important; }
.sidebar-header { padding: 24px !important; border-bottom: 1px solid var(--border-color) !important; }
.logo { display: flex !important; align-items: center !important; gap: 12px !important; font-weight: 600 !important; font-size: 16px !important; letter-spacing: -0.02em !important; color: #171717 !important; }
.logo svg { stroke: #171717 !important; }
.sidebar-content { flex-grow: 1 !important; padding: 24px !important; display: flex !important; flex-direction: column !important; gap: 20px !important; }
.input-group { display: flex !important; flex-direction: column !important; gap: 8px !important; flex-grow: 1 !important; }
.input-group label { font-size: 13px !important; font-weight: 500 !important; color: var(--text-secondary) !important; }

textarea.custom-textarea {
    flex-grow: 1 !important; width: 100% !important; height: 100% !important; resize: none !important; border: 1px solid var(--border-color) !important; border-radius: 8px !important; padding: 16px !important; font-family: inherit !important; font-size: 14px !important; line-height: 1.6 !important; color: var(--text-primary) !important; background-color: #fff !important; transition: all 0.2s ease !important; box-shadow: 0 1px 2px rgba(0,0,0,0.02) !important;
}
textarea.custom-textarea:focus { outline: none !important; border-color: #999 !important; box-shadow: 0 0 0 4px var(--focus-ring) !important; }
textarea.custom-textarea::placeholder { color: #a1a1aa !important; }

.btn-primary { background-color: #000 !important; color: #fff !important; border: none !important; border-radius: 6px !important; padding: 12px 16px !important; font-size: 14px !important; font-weight: 500 !important; cursor: pointer !important; transition: all 0.2s ease !important; display: flex !important; align-items: center !important; justify-content: center !important; gap: 8px !important; }
.btn-primary, .btn-primary * { color: #ffffff !important; stroke: #ffffff !important; }
.btn-primary:hover { background-color: #333 !important; }
.btn-primary:active { transform: scale(0.98) !important; }
.btn-primary:disabled { background-color: #e5e5e5 !important; color: #a3a3a3 !important; cursor: not-allowed !important; stroke: #a3a3a3 !important; }

#loading { display: flex; align-items: center !important; justify-content: center !important; gap: 12px !important; font-size: 13px !important; color: var(--text-secondary) !important; padding: 12px !important; }
.spinner { width: 16px !important; height: 16px !important; border: 2px solid var(--border-color) !important; border-top: 2px solid #000 !important; border-radius: 50% !important; animation: spin 0.8s linear infinite !important; }
@keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } }
#loading.hidden { display: none !important; }

.sidebar-footer { padding: 16px 24px !important; border-top: 1px solid var(--border-color) !important; font-size: 12px !important; color: var(--text-tertiary) !important; }
.sidebar-footer p { color: var(--text-tertiary) !important; }

.canvas-area { flex-grow: 1 !important; position: relative !important; background-color: var(--bg-canvas) !important; background-image: radial-gradient(#e5e7eb 1px, transparent 1px) !important; background-size: 20px 20px !important; padding: 0 !important; margin: 0 !important; height: 100vh !important; }
#markmap { width: 100% !important; height: 100% !important; }
#markmap text { fill: #171717 !important; color: #171717 !important; font-family: inherit !important; }

.app-layout .disclaimer { color: #a1a1aa !important; }

/* Gradio Overrides to remove padding and ensure full screen */
.gradio-container { padding: 0 !important; margin: 0 !important; max-width: 100vw !important; border: none !important; }
footer { display: none !important; }
#hidden-layer { display: none !important; }
"""

original_html = """
<div class="app-layout">
    <aside class="sidebar">
        <div class="sidebar-header">
            <div class="logo">
                <svg width="24" height="24" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 16V8a2 2 0 0 0-1-1.73l-7-4a2 2 0 0 0-2 0l-7 4A2 2 0 0 0 3 8v8a2 2 0 0 0 1 1.73l7 4a2 2 0 0 0 2 0l7-4A2 2 0 0 0 21 16z"></path><polyline points="3.27 6.96 12 12.01 20.73 6.96"></polyline><line x1="12" y1="22.08" x2="12" y2="12"></line></svg>
                <span>MindMap Studio</span>
            </div>
        </div>
        <div class="sidebar-content">
            <div class="input-group">
                <label for="inputText">Source Text</label>
                <textarea id="inputText" class="custom-textarea" placeholder="議事録や講義のテキストをペーストしてください..."></textarea>
            </div>
            <button id="generateBtn" class="btn-primary">
                <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polygon points="5 3 19 12 5 21 5 3"></polygon></svg>
                マップを生成
            </button>
            <div id="loading" class="hidden">
                <div class="spinner"></div>
                <span>Processing text...</span>
            </div>
        </div>
        <div class="sidebar-footer">
            <p>Powered by edha 1.0 3B</p>
        </div>
    </aside>
    <main class="canvas-area">
        <svg id="markmap"></svg>
        <div class="disclaimer" style="position: absolute; bottom: 16px; left: 50%; transform: translateX(-50%); font-size: 11px; color: #a1a1aa; text-align: center; pointer-events: none; z-index: 1000; width: 100%;">
            MindMap Studioの回答は正しいとは限らないので、重要な情報は必ず見直してください。
        </div>
    </main>
</div>
"""

head_scripts = """
<script src="https://cdn.jsdelivr.net/npm/d3@7"></script>
<script src="https://cdn.jsdelivr.net/npm/markmap-lib"></script>
<script src="https://cdn.jsdelivr.net/npm/markmap-view"></script>
<script>
document.addEventListener('DOMContentLoaded', () => {
    const { markmap } = window;
    const { Markmap, loadCSS, loadJS, Transformer } = markmap;
    const transformer = new Transformer();
    let mm = null;

    // Gradio injects elements dynamically, so we wait until our HTML and Gradio's hidden inputs are ready
    const initInterval = setInterval(() => {
        const svgEl = document.querySelector('#markmap');
        const hiddenBtn = document.querySelector('#hidden-btn');
        if (svgEl && hiddenBtn && !mm) {
            clearInterval(initInterval);
            mm = Markmap.create('#markmap');
            
            const initialMarkdown = `
# マインドマップ自動生成
## 使い方
- 左側に文章を入力します
- 「マップを生成」ボタンを押します
## 特徴
- AIが文脈を理解して自動で構造化
- 専用カスタムAI（edha 1.0 3B）による情報抽出
`;
            renderMindMap(initialMarkdown);
            setupBridge();
        }
    }, 100);

    function renderMindMap(markdownContent) {
        const { root, features } = transformer.transform(markdownContent);
        const { styles, scripts } = transformer.getUsedAssets(features);
        if (styles) loadCSS(styles);
        if (scripts) loadJS(scripts, { getMarkmap: () => markmap });
        mm.setData(root);
        mm.fit();
    }

    function setupBridge() {
        const generateBtn = document.getElementById('generateBtn');
        const inputText = document.getElementById('inputText');
        const loadingDiv = document.getElementById('loading');
        
        generateBtn.addEventListener('click', () => {
            const text = inputText.value.trim();
            if (!text) return;
            
            generateBtn.disabled = true;
            generateBtn.classList.add('hidden');
            loadingDiv.classList.remove('hidden');
            
            // Bridge custom textarea to Gradio's hidden textarea
            const hiddenInput = document.querySelector('#hidden-input textarea');
            const hiddenBtn = document.querySelector('#hidden-btn');
            
            if (hiddenInput && hiddenBtn) {
                hiddenInput.value = text;
                hiddenInput.dispatchEvent(new Event('input', { bubbles: true }));
                setTimeout(() => hiddenBtn.click(), 50); // let Svelte process input
            }
        });
        
        // Listen to changes on Gradio's hidden output
        const hiddenOutput = document.querySelector('#hidden-output textarea');
        if (hiddenOutput) {
            let lastVal = hiddenOutput.value;
            setInterval(() => {
                if (hiddenOutput.value !== lastVal) {
                    lastVal = hiddenOutput.value;
                    if (lastVal && lastVal.trim().length > 0) {
                        renderMindMap(lastVal);
                        
                        generateBtn.disabled = false;
                        generateBtn.classList.remove('hidden');
                        loadingDiv.classList.add('hidden');
                    }
                }
            }, 200);
        }
    }
});
</script>
"""

with gr.Blocks(css=base_css, head=head_scripts, title="MindMap Studio") as demo:
    # 完全にオリジナルのUIをインジェクト（Gradioのコンテナ制限を受けない純粋なHTML文字列）
    gr.HTML(original_html)
    
    # Python関数の実行に必要なGradioコンポーネント（目に見えないように隠蔽）
    with gr.Row(elem_id="hidden-layer"):
        hidden_input = gr.Textbox(elem_id="hidden-input")
        hidden_output = gr.Textbox(elem_id="hidden-output")
        hidden_btn = gr.Button(elem_id="hidden-btn")
        
    hidden_btn.click(fn=generate_mindmap, inputs=hidden_input, outputs=hidden_output, api_name="generate")

demo.launch()