#!/usr/bin/env python3
"""Apply v1.4.0 improvements: model color fix, review button, web search grounding."""

import re

fp = '/home/uroma/promptarch/components/AIAssist.tsx'
with open(fp, 'r') as f:
    content = f.read()

changes = 0

# ============================================================
# CHANGE 1: Fix model selector text color
# ============================================================
old_option = '{availableModels.map(m => <option key={m} value={m}>{m}</option>)}'
new_option = '{availableModels.map(m => <option key={m} value={m} className="bg-[#0b1414] text-white">{m}</option>)}'

if old_option in content:
    content = content.replace(old_option, new_option, 1)
    changes += 1
    print("Change 1: Fixed <option> text color to white")
else:
    print("WARNING: Could not find option map")

# Also add text-white to the <select> className
old_select_class = 'className="text-[10px] font-bold h-8 px-2 rounded-lg border border-blue-100 dark:border-blue-900 bg-white/80 dark:bg-[#0b1414]/80 focus:ring-2 focus:ring-blue-400/30 transition-all outline-none min-w-[120px]"'
new_select_class = 'className="text-[10px] font-bold h-8 px-2 rounded-lg border border-blue-100 dark:border-blue-900 bg-white/80 dark:bg-[#0b1414]/80 focus:ring-2 focus:ring-blue-400/30 transition-all outline-none min-w-[120px] text-slate-900 dark:text-white"'

if old_select_class in content:
    content = content.replace(old_select_class, new_select_class, 1)
    changes += 1
    print("Change 1b: Added text-white to select element")
else:
    print("WARNING: Could not find select className")

# ============================================================
# CHANGE 2: Add ShieldCheck to lucide-react imports
# ============================================================
old_import = '    Wand2, LayoutPanelLeft, Play, Orbit, Plus, Key'
new_import = '    Wand2, LayoutPanelLeft, Play, Orbit, Plus, Key, ShieldCheck'

if old_import in content:
    content = content.replace(old_import, new_import, 1)
    changes += 1
    print("Change 2: Added ShieldCheck to imports")
else:
    print("WARNING: Could not find lucide import line")

# ============================================================
# CHANGE 3: Add reviewCode function after approveAndGenerate
# ============================================================
old_approve = '''    const approveAndGenerate = () => {
        setAssistStep("generating");
        handleSendMessage(undefined, "Approved. Please generate the code according to the plan.", true);
    };'''

new_approve = '''    const approveAndGenerate = () => {
        setAssistStep("generating");
        handleSendMessage(undefined, "Approved. Please generate the code according to the plan.", true);
    };

    const reviewCode = () => {
        if (!previewData?.data) return;
        setAssistStep("generating");
        const reviewPrompt = `Please review this generated code for bugs, security issues, performance problems, and best practices. Provide specific improvements:\\n\\n\\\`\\\`\\\`${previewData.language || "code"}\\n${previewData.data}\\n\\\`\\\`\\\``;
        handleSendMessage(undefined, reviewPrompt, true);
    };'''

if old_approve in content:
    content = content.replace(old_approve, new_approve, 1)
    changes += 1
    print("Change 3: Added reviewCode function")
else:
    print("WARNING: Could not find approveAndGenerate block")

# ============================================================
# CHANGE 4: Add Review Code button (grid-cols-2 -> 3)
# ============================================================
old_grid = '''<div className="mt-4 grid grid-cols-2 gap-2 animate-in zoom-in-95 duration-300">
                                            <Button
                                                onClick={() => { setShowCanvas(true); setViewMode(isPreviewRenderable(previewData as PreviewData) ? "preview" : "code"); }}
                                                className="bg-blue-600 hover:bg-blue-500 text-white font-black uppercase text-[10px] tracking-widest py-4 rounded-xl shadow-lg shadow-blue-500/20"
                                            >
                                                <Zap className="h-3.5 w-3.5 mr-1.5" /> {t.activateArtifact}
                                            </Button>
                                            <Button
                                                onClick={() => { setAssistStep("idle"); setInput("Modify this: "); setTimeout(() => { const el = document.querySelector<HTMLInputElement>(`[data-ai-input]`); if (el) el.focus(); }, 100); }}
                                                variant="outline"
                                                className="bg-slate-500/10 hover:bg-slate-500/20 border-slate-500/20 text-slate-300 font-black uppercase text-[10px] tracking-widest py-4 rounded-xl"
                                            >
                                                <LayoutPanelLeft className="h-3.5 w-3.5 mr-1.5" /> Request Modifications
                                            </Button>
                                        </div>'''

new_grid = '''<div className="mt-4 grid grid-cols-3 gap-2 animate-in zoom-in-95 duration-300">
                                            <Button
                                                onClick={() => { setShowCanvas(true); setViewMode(isPreviewRenderable(previewData as PreviewData) ? "preview" : "code"); }}
                                                className="bg-blue-600 hover:bg-blue-500 text-white font-black uppercase text-[10px] tracking-widest py-4 rounded-xl shadow-lg shadow-blue-500/20"
                                            >
                                                <Zap className="h-3.5 w-3.5 mr-1.5" /> {t.activateArtifact}
                                            </Button>
                                            <Button
                                                onClick={reviewCode}
                                                disabled={!previewData?.data}
                                                variant="outline"
                                                className="bg-emerald-500/10 hover:bg-emerald-500/20 border-emerald-500/20 text-emerald-300 font-black uppercase text-[10px] tracking-widest py-4 rounded-xl"
                                            >
                                                <ShieldCheck className="h-3.5 w-3.5 mr-1.5" /> Review Code
                                            </Button>
                                            <Button
                                                onClick={() => { setAssistStep("idle"); setInput("Modify this: "); setTimeout(() => { const el = document.querySelector<HTMLInputElement>(`[data-ai-input]`); if (el) el.focus(); }, 100); }}
                                                variant="outline"
                                                className="bg-slate-500/10 hover:bg-slate-500/20 border-slate-500/20 text-slate-300 font-black uppercase text-[10px] tracking-widest py-4 rounded-xl"
                                            >
                                                <LayoutPanelLeft className="h-3.5 w-3.5 mr-1.5" /> Request Modifications
                                            </Button>
                                        </div>'''

if old_grid in content:
    content = content.replace(old_grid, new_grid, 1)
    changes += 1
    print("Change 4: Added Review Code button (grid-cols-3)")
else:
    print("WARNING: Could not find post-coding buttons block")

# ============================================================
# CHANGE 5: Add webSearchEnabled state after isProcessing
# ============================================================
old_state = '    const [isProcessing, setIsProcessing] = useState(false);'
new_state = '    const [isProcessing, setIsProcessing] = useState(false);\n    const [webSearchEnabled, setWebSearchEnabled] = useState(false);'

if old_state in content:
    content = content.replace(old_state, new_state, 1)
    changes += 1
    print("Change 5: Added webSearchEnabled state")
else:
    print("WARNING: Could not find isProcessing state")

# ============================================================
# CHANGE 6: Add web search toggle button in toolbar
# (After the model select closing </select> and before the
#  canvas toggle div)
# ============================================================
old_toolbar = '''                                </select>
                                <div className="flex items-center gap-1 ml-1 border-l border-blue-100 dark:border-blue-900 pl-2">'''

new_toolbar = '''                                </select>
                                <button
                                    onClick={() => setWebSearchEnabled(prev => !prev)}
                                    className={`flex items-center gap-1 h-8 px-2 rounded-lg text-[10px] font-bold transition-all border ${
                                        webSearchEnabled
                                            ? "bg-amber-500/20 border-amber-500/40 text-amber-300"
                                            : "bg-transparent border-transparent text-slate-400/70 hover:text-slate-300 hover:bg-slate-500/10"
                                    }`}
                                >
                                    <Search className="h-3.5 w-3.5" /> {webSearchEnabled ? "ON" : "Search"}
                                </button>
                                <div className="flex items-center gap-1 ml-1 border-l border-blue-100 dark:border-blue-900 pl-2">'''

if old_toolbar in content:
    content = content.replace(old_toolbar, new_toolbar, 1)
    changes += 1
    print("Change 6: Added web search toggle button")
else:
    print("WARNING: Could not find toolbar insertion point")

# ============================================================
# CHANGE 7: Add web search enrichment in handleSendMessage
# Before generateAIAssistStream call, after formattedHistory
# ============================================================
old_stream = '''            const response = await modelAdapter.generateAIAssistStream(
                {
                    messages: [...formattedHistory, { role: "user" as const, content: finalInput, timestamp: new Date() }],'''

new_stream = '''            // Web search grounding: enrich prompt with search results if enabled
            let enrichedInput = finalInput;
            if (webSearchEnabled) {
                setStatus("Searching the web...");
                try {
                    const searchRes = await fetch("/api/search?q=" + encodeURIComponent(finalInput.split("\\n")[0].substring(0, 200)));
                    if (searchRes.ok) {
                        const searchData = await searchRes.json();
                        if (searchData.results && searchData.results.length > 0) {
                            const searchContext = searchData.results.slice(0, 5).map((r: { title: string; url: string; snippet: string }, i: number) =>
                                `${i + 1}. **${r.title}** (${r.url}) - ${r.snippet}`
                            ).join("\\n");
                            enrichedInput = `[WEB SEARCH CONTEXT - Top 5 relevant results]\\n${searchContext}\\n\\n---\\nUsing the above search results as reference context, answer the user query. Cite sources when relevant.\\n\\nUser query: ${finalInput}`;
                        }
                    }
                } catch (e) {
                    console.warn("Web search failed:", e);
                }
                setStatus(null);
            }

            const response = await modelAdapter.generateAIAssistStream(
                {
                    messages: [...formattedHistory, { role: "user" as const, content: enrichedInput, timestamp: new Date() }],'''

if old_stream in content:
    content = content.replace(old_stream, new_stream, 1)
    changes += 1
    print("Change 7: Added web search enrichment before AI call")
else:
    print("WARNING: Could not find generateAIAssistStream call")

with open(fp, 'w') as f:
    f.write(content)

print(f"\nApplied {changes}/7 changes to AIAssist.tsx")