# Copyright Volkan Sah aka AliBrown
# Free to use with your own api-key
import os
import streamlit as st
import tempfile
import requests
import json
from datetime import datetime

# ----------------------------------------------------
# 🚨 HUGGINGFACE SPACES FIX
# ----------------------------------------------------
TEMP_STREAMLIT_HOME = os.path.join(tempfile.gettempdir(), "st_config_workaround")
os.makedirs(TEMP_STREAMLIT_HOME, exist_ok=True)

os.environ["STREAMLIT_HOME"] = TEMP_STREAMLIT_HOME
os.environ["STREAMLIT_GATHER_USAGE_STATS"] = "false"

CONFIG_PATH = os.path.join(TEMP_STREAMLIT_HOME, "config.toml")
CONFIG_CONTENT = """
[browser]
gatherUsageStats = false
"""

if not os.path.exists(CONFIG_PATH):
    try:
        with open(CONFIG_PATH, "w") as f:
            f.write(CONFIG_CONTENT)
    except Exception:
        pass

# ----------------------------------------------------
# CONFIG
# ----------------------------------------------------
st.set_page_config(
    page_title="AI Code Forge", 
    layout="wide",
    page_icon="⚒️"
)

OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"

FREE_MODELS = {
    "🧠 DeepSeek Chat": "deepseek/deepseek-chat-v3.1:free",
    "💻 Qwen Coder": "qwen/qwen3-coder:free",
    "🔮 Gemma 27B": "google/gemma-3-27b-it:free",
    "🐬 Dolphin Mistral": "cognitivecomputations/dolphin-mistral-24b-venice-edition:free",
    "⚡ Nemotron Nano": "nvidia/nemotron-nano-9b-v2:free",
}

# ----------------------------------------------------
# SESSION STATE
# ----------------------------------------------------
if "messages" not in st.session_state:
    st.session_state.messages = {}
if "current_tool" not in st.session_state:
    st.session_state.current_tool = "home"
if "last_output" not in st.session_state:
    st.session_state.last_output = ""
if "dark_mode" not in st.session_state:
    st.session_state.dark_mode = True

# ----------------------------------------------------
# HELPER FUNCTIONS
# ----------------------------------------------------
def fetch_model_contexts(api_key):
    if not api_key:
        return {}
    headers = {"Authorization": f"Bearer {api_key}"}
    try:
        res = requests.get(f"{OPENROUTER_API_BASE}/models", headers=headers, timeout=10)
        contexts = {}
        if res.status_code == 200:
            for m in res.json().get("data", []):
                contexts[m.get("id")] = m.get("context_length", 4096)
        return contexts
    except Exception:
        return {}

def call_openrouter(model, messages, temp, max_tok, key, system_prompt=None):
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "Referer": "https://aicodecraft.io",
        "X-Title": "AI-Code-Forge",
    }
    
    api_messages = messages.copy()
    if system_prompt:
        api_messages.insert(0, {"role": "system", "content": system_prompt})
    
    payload = {
        "model": model,
        "messages": api_messages,
        "temperature": temp,
        "max_tokens": max_tok,
    }
    
    res = requests.post(
        f"{OPENROUTER_API_BASE}/chat/completions", 
        headers=headers, 
        data=json.dumps(payload),
        timeout=60
    )
    
    if res.status_code == 200:
        try:
            return res.json()["choices"][0]["message"]["content"]
        except (KeyError, IndexError):
            raise Exception("Fehlerhafte API-Antwort")
    else:
        try:
            err = res.json()
            msg = err.get("error", {}).get("message", res.text)
        except:
            msg = res.text
        raise Exception(f"API Error {res.status_code}: {msg}")

def extract_code_blocks(text):
    """Extrahiert Code aus Markdown Code-Blöcken"""
    import re
    pattern = r"```(\w+)?\n(.*?)```"
    matches = re.findall(pattern, text, re.DOTALL)
    if matches:
        return [(lang or "text", code.strip()) for lang, code in matches]
    return [("text", text)]

# ----------------------------------------------------
# TOOL PRESETS
# ----------------------------------------------------
TOOL_PRESETS = {
    "code_gen": {
        "icon": "🎨",
        "title": "Code Generator",
        "color": "#00ff88",
        "system_prompt": "You are an Expert Code Generator. Create clean, production-ready code with comments. Respond ONLY with code in Markdown code blocks.",
        "templates": {
            "Flask REST API": "Create a Flask REST API with:\n- CRUD Endpoints for 'tasks'\n- SQLAlchemy Models\n- Error Handling\n- CORS enabled",
            "React Component": "Create a React Functional Component with:\n- Props for title and items (Array)\n- useState for selected item\n- Tailwind CSS Styling\n- Click Handler",
            "Python CLI Tool": "Create a Python CLI Tool with:\n- argparse for Arguments\n- Logging\n- Error Handling\n- Main Guard",
            "FastAPI Endpoint": "Create FastAPI Endpoints:\n- POST /users - Create User\n- GET /users/{id} - Get User\n- Pydantic Models\n- SQLAlchemy Integration",
        }
    },
    "refactor": {
        "icon": "🔄",
        "title": "Code Refactorer",
        "color": "#00aaff",
        "system_prompt": "You are a Code Refactoring Expert. Analyze code and improve: readability, performance, best practices. Explain your changes briefly.",
        "templates": {
            "Clean Code": "Refactor this code according to Clean Code principles:\n- Improve Naming\n- Reduce Complexity\n- Extract Functions\n\n[Insert Code Here]",
            "Performance": "Optimize this code for performance:\n- Identify Bottlenecks\n- Improve Algorithms\n- Reduce Memory Usage\n\n[Insert Code Here]",
            "Type Safety": "Add Type Hints/TypeScript:\n- All Functions\n- Variables where appropriate\n- Return Types\n\n[Insert Code Here]",
        }
    },
    "debug": {
        "icon": "🐛",
        "title": "Debug Helper",
        "color": "#ff4444",
        "system_prompt": "You are a Debugging Expert. Analyze code/errors systematically. Explain the problem and provide the fix.",
        "templates": {
            "Error Analysis": "I have this error:\n[Error Message]\n\nIn this code:\n[Code]\n\nWhat is the problem and how do I fix it?",
            "Logic Bug": "This code doesn't do what it's supposed to:\n[Code]\n\nExpected behavior: [describe]\nActual behavior: [describe]\n\nWhere is the logic bug?",
            "Performance Issue": "This code is too slow:\n[Code]\n\nInput Size: [describe]\nCurrent Runtime: [describe]\n\nWhere is the problem?",
        }
    },
    "docs": {
        "icon": "📚",
        "title": "Doc Generator",
        "color": "#ffaa00",
        "system_prompt": "You are a Technical Writer. Create clear, structured documentation with examples.",
        "templates": {
            "Function Docs": "Create Docstrings for these functions:\n[Code]\n\nFormat: Google Style for Python / JSDoc for JavaScript",
            "README": "Create a README.md for this project:\n- Description\n- Installation\n- Usage Examples\n- API Reference\n\nProject: [describe]",
            "API Docs": "Create API Documentation for these Endpoints:\n[Code/Routes]\n\nIncl. Request/Response Examples",
        }
    },
    "explain": {
        "icon": "💡",
        "title": "Code Explainer",
        "color": "#aa88ff",
        "system_prompt": "You are a Code Educator. Explain code step-by-step, clearly and understandably.",
        "templates": {
            "ELI5": "Explain this code as if I were 5:\n[Code]",
            "Deep Dive": "Explain this code in detail:\n- What does it do?\n- How does it work?\n- Why was it solved this way?\n- Alternatives?\n\n[Code]",
            "Architecture": "Explain the architecture of this code:\n- Design Patterns\n- Components\n- Data Flow\n\n[Code]",
        }
    }
}
# ----------------------------------------------------
# SIDEBAR
# ----------------------------------------------------
with st.sidebar:
    st.markdown("# ⚒️ AI Code Forge")
    st.markdown("---")
    
    # API Settings
    with st.expander("🔑 API Settings", expanded=True):
        api_key = st.text_input("OpenRouter API Key", type="password")
        
        model_name = st.selectbox("Model", list(FREE_MODELS.keys()))
        model = FREE_MODELS[model_name]
        
        model_contexts = fetch_model_contexts(api_key)
        default_ctx = model_contexts.get(model, 4096)
        
        temperature = st.slider("Temperature", 0.0, 1.0, 0.7, 0.1)
        max_tokens = st.slider(
            "Max Tokens", 
            256, 
            min(default_ctx, 16000), 
            min(2048, default_ctx),
            step=256
        )
        
        st.caption(f"📊 Context: {default_ctx:,} tokens")
    
    st.markdown("---")
    
    # Navigation
    st.markdown("### 🧭 Tools")
    
    if st.button("🏠 Home", use_container_width=True):
        st.session_state.current_tool = "home"
        st.rerun()
    
    for tool_id, tool in TOOL_PRESETS.items():
        if st.button(
            f"{tool['icon']} {tool['title']}", 
            use_container_width=True,
            type="primary" if st.session_state.current_tool == tool_id else "secondary"
        ):
            st.session_state.current_tool = tool_id
            if tool_id not in st.session_state.messages:
                st.session_state.messages[tool_id] = []
            st.rerun()
    
    st.markdown("---")
    
    # Actions
    if st.session_state.current_tool != "home":
        if st.button("🗑️ Clear Chat", use_container_width=True):
            if st.session_state.current_tool in st.session_state.messages:
                st.session_state.messages[st.session_state.current_tool] = []
            st.success("Chat cleared!")
            st.rerun()
    
    st.markdown("---")
    st.caption("💸 Using Free Models")
    st.caption(f"⏰ {datetime.now().strftime('%H:%M')}")

# ----------------------------------------------------
# MAIN CONTENT
# ----------------------------------------------------

# HOME
if st.session_state.current_tool == "home":
    st.markdown("# ⚒️ Welcome to AI Code Forge")
    st.markdown("### Your AI-Powered Development Assistant")
    
    st.markdown("---")
    
    # Tool Cards
    cols = st.columns(3)
    tool_items = list(TOOL_PRESETS.items())
    
    for idx, (tool_id, tool) in enumerate(tool_items):
        with cols[idx % 3]:
            with st.container():
                st.markdown(f"""
                <div style="
                    padding: 20px;
                    border-radius: 10px;
                    border-left: 4px solid {tool['color']};
                    background: rgba(255,255,255,0.05);
                    margin-bottom: 20px;
                ">
                    <h2>{tool['icon']} {tool['title']}</h2>
                    <p style="opacity: 0.8;">{tool['system_prompt'][:100]}...</p>
                </div>
                """, unsafe_allow_html=True)
                
                if st.button(f"Open {tool['title']}", key=f"open_{tool_id}", use_container_width=True):
                    st.session_state.current_tool = tool_id
                    if tool_id not in st.session_state.messages:
                        st.session_state.messages[tool_id] = []
                    st.rerun()
    
    st.markdown("---")
    
    # Quick Stats
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Free Models", len(FREE_MODELS))
    with col2:
        total_msgs = sum(len(msgs) for msgs in st.session_state.messages.values())
        st.metric("Total Messages", total_msgs)
    with col3:
        st.metric("Tools", len(TOOL_PRESETS))

# TOOL VIEW
else:
    tool = TOOL_PRESETS[st.session_state.current_tool]
    
    # Header
    st.markdown(f"""
    <div style="
        padding: 20px;
        border-radius: 10px;
        border-left: 6px solid {tool['color']};
        background: rgba(255,255,255,0.05);
        margin-bottom: 30px;
    ">
        <h1>{tool['icon']} {tool['title']}</h1>
        <p style="opacity: 0.8; font-size: 1.1em;">{tool['system_prompt']}</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Templates
    with st.expander("📋 Quick Templates", expanded=False):
        for template_name, template_text in tool['templates'].items():
            if st.button(f"📄 {template_name}", key=f"template_{template_name}"):
                st.session_state[f"template_text_{st.session_state.current_tool}"] = template_text
                st.rerun()
    
    # Chat Interface
    st.markdown("### 💬 Chat")
    
    # Display messages
    for msg in st.session_state.messages.get(st.session_state.current_tool, []):
        with st.chat_message(msg["role"]):
            if msg["role"] == "assistant":
                # Check for code blocks
                code_blocks = extract_code_blocks(msg["content"])
                if len(code_blocks) == 1 and code_blocks[0][0] == "text":
                    st.markdown(msg["content"])
                else:
                    # Has code blocks
                    parts = msg["content"].split("```")
                    for i, part in enumerate(parts):
                        if i % 2 == 0:
                            # Text part
                            if part.strip():
                                st.markdown(part)
                        else:
                            # Code part
                            lines = part.split("\n", 1)
                            lang = lines[0].strip() if lines else "text"
                            code = lines[1] if len(lines) > 1 else part
                            
                            col1, col2 = st.columns([6, 1])
                            with col1:
                                st.code(code, language=lang)
                            with col2:
                                if st.button("📋", key=f"copy_{i}_{msg.get('timestamp', 0)}"):
                                    st.toast("Code copied! (simulation)")
            else:
                st.markdown(msg["content"])
    
    # Input
    default_text = st.session_state.get(f"template_text_{st.session_state.current_tool}", "")
    if default_text:
        del st.session_state[f"template_text_{st.session_state.current_tool}"]
    
    if prompt := st.chat_input("Your message...", key=f"input_{st.session_state.current_tool}"):
        user_prompt = prompt
    elif default_text:
        user_prompt = default_text
    else:
        user_prompt = None
    
    if user_prompt:
        if not api_key:
            st.warning("⚠️ Please enter your OpenRouter API Key in the sidebar!")
            st.stop()
        
        # Add user message
        st.session_state.messages[st.session_state.current_tool].append({
            "role": "user", 
            "content": user_prompt,
            "timestamp": datetime.now().timestamp()
        })
        
        with st.chat_message("user"):
            st.markdown(user_prompt)
        
        # Prepare messages for API
        api_messages = [
            {"role": m["role"], "content": m["content"]} 
            for m in st.session_state.messages[st.session_state.current_tool]
        ]
        
        # Generate response
        with st.chat_message("assistant"):
            with st.spinner(f"🤖 {model_name} thinking..."):
                try:
                    reply = call_openrouter(
                        model, 
                        api_messages, 
                        temperature, 
                        max_tokens, 
                        api_key,
                        system_prompt=tool['system_prompt']
                    )
                    
                    # Display with code highlighting
                    code_blocks = extract_code_blocks(reply)
                    if len(code_blocks) == 1 and code_blocks[0][0] == "text":
                        st.markdown(reply)
                    else:
                        parts = reply.split("```")
                        for i, part in enumerate(parts):
                            if i % 2 == 0:
                                if part.strip():
                                    st.markdown(part)
                            else:
                                lines = part.split("\n", 1)
                                lang = lines[0].strip() if lines else "text"
                                code = lines[1] if len(lines) > 1 else part
                                st.code(code, language=lang)
                    
                    st.session_state.messages[st.session_state.current_tool].append({
                        "role": "assistant", 
                        "content": reply,
                        "timestamp": datetime.now().timestamp()
                    })
                    
                except Exception as e:
                    error_msg = f"❌ Error: {str(e)}"
                    st.error(error_msg)
                    st.session_state.messages[st.session_state.current_tool].append({
                        "role": "assistant", 
                        "content": error_msg,
                        "timestamp": datetime.now().timestamp()
                    })