import os
import time

# Check if spaces library is available (for Hugging Face Spaces ZeroGPU)
# Importing spaces BEFORE torch is a strict ZeroGPU requirement
try:
    import spaces
    has_spaces = True
    print("Hugging Face Spaces library loaded successfully.")
except ImportError:
    has_spaces = False
    print("Hugging Face Spaces library not found. Running in standard environment.")

import torch
from typing import Generator
from fastapi.responses import HTMLResponse
from gradio import Server

# Initialize the Gradio Server (extends FastAPI)
app = Server()

# Define HTML path
HTML_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")

# Set device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"System detected device: {device}")

# Determine if model should be loaded
force_load = os.environ.get("FORCE_MODEL_LOAD", "false").lower() == "true"
is_fallback = True
model = None
tokenizer = None

MODEL_ID = "KyleHessling1/Qwopus3.6-27B-Fusion-GGUF:Q4_K_M"

if device == "cuda" or force_load:
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, GenerationConfig
        from threading import Thread
        
        print(f"Attempting to load model '{MODEL_ID}'...")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        
        if torch.backends.mps.is_available() and force_load:
            print("Loading model on MPS (Apple Silicon GPU) with float16...")
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_ID,
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True
            ).to("mps")
        else:
            print("Loading model in bfloat16...")
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_ID,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                low_cpu_mem_usage=True
            )
        
        is_fallback = False
        print("Model loaded successfully!")
    except Exception as e:
        print(f"Error loading model: {e}")
        print("Falling back to simulation mode.")
        is_fallback = True
else:
    print("No CUDA GPU detected and FORCE_MODEL_LOAD is false. Falling back to simulation mode.")
    is_fallback = True


def get_mock_response(message: str) -> tuple[str, str]:
    """Generates detailed mock thinking and answers for local development testing."""
    message_lower = message.lower()
    
    if "palindrom" in message_lower or "leetcode" in message_lower or "code" in message_lower or "python" in message_lower:
        thought = (
            "1. We need to solve the Longest Palindromic Substring problem.\n"
            "2. Let's analyze the constraints and possible approaches.\n"
            "   - Approach 1: Brute Force. Check all O(N^2) substrings. Checking takes O(N), total O(N^3). Too slow.\n"
            "   - Approach 2: Dynamic Programming. Let DP[i][j] be true if substring s[i..j] is a palindrome.\n"
            "     - DP[i][j] = (s[i] == s[j]) && (j - i < 3 || DP[i+1][j-1])\n"
            "     - Time complexity: O(N^2), Space complexity: O(N^2).\n"
            "   - Approach 3: Expand Around Center. For each index, expand outward for both odd and even length palindromes.\n"
            "     - Time complexity: O(N^2), Space complexity: O(1). This is optimal in terms of space.\n"
            "   - Approach 4: Manacher's Algorithm. Dynamic programming combined with centers expansion. O(N) time and space.\n"
            "3. Let's implement the Expand Around Center approach as it is highly readable and O(1) space.\n"
            "4. Verification: check boundary cases like single character, empty string, string with all identical characters, etc.\n"
            "5. Formulate final response with explanation, code, and complexity analysis."
        )
        body = (
            "Here is the optimal Python implementation of the **Longest Palindromic Substring** problem using the **Expand Around Center** approach (\\(O(N^2)\\) time, \\(O(1)\\) space).\n\n"
            "### Expand Around Center (Python)\n\n"
            "```python\n"
            "class Solution:\n"
            "    def longestPalindrome(self, s: str) -> str:\n"
            "        if not s or len(s) < 1:\n"
            "            return \"\"\n"
            "        \n"
            "        start, end = 0, 0\n"
            "        \n"
            "        def expand_around_center(left: int, right: int) -> int:\n"
            "            while left >= 0 and right < len(s) and s[left] == s[right]:\n"
            "                left -= 1\n"
            "                right += 1\n"
            "            # Return the length of the palindrome found\n"
            "            return right - left - 1\n"
            "        \n"
            "        for i in range(len(s)):\n"
            "            # Odd-length palindromes (single character center)\n"
            "            len1 = expand_around_center(i, i)\n"
            "            # Even-length palindromes (two character center)\n"
            "            len2 = expand_around_center(i, i + 1)\n"
            "            \n"
            "            max_len = max(len1, len2)\n"
            "            if max_len > end - start:\n"
            "                # Adjust start and end indices based on current center\n"
            "                start = i - (max_len - 1) // 2\n"
            "                end = i + max_len // 2\n"
            "                \n"
            "        return s[start:end + 1]\n"
            "```\n\n"
            "### Complexity Analysis\n"
            "- **Time Complexity:** \\(O(N^2)\\). We expand around \\(2N - 1\\) centers. For each center, expansion can take up to \\(O(N)\\) steps.\n"
            "- **Space Complexity:** \\(O(1)\\). Only constant extra space is used."
        )
    elif "deck" in message_lower or "probab" in message_lower or "card" in message_lower or "math" in message_lower or "solve" in message_lower or "equation" in message_lower:
        thought = (
            "1. The user is asking a probability/math question: 'If a card is drawn from a standard deck, what is the probability that it is a spade or a face card?'\n"
            "2. Let's define the sample space and events:\n"
            "   - Total cards in a standard deck: N(S) = 52.\n"
            "   - Event A: Drawing a spade. There are 13 spades in a deck. So N(A) = 13.\n"
            "   - Event B: Drawing a face card (Jack, Queen, King). There are 3 face cards per suit, and 4 suits. So N(B) = 3 * 4 = 12.\n"
            "   - We need to find the probability of Spade OR Face Card: P(A or B).\n"
            "3. Let's recall the addition rule of probability:\n"
            "   - P(A or B) = P(A) + P(B) - P(A and B)\n"
            "4. What is Event (A and B)? It is drawing a card that is both a spade AND a face card.\n"
            "   - These are the Jack of Spades, Queen of Spades, and King of Spades. N(A and B) = 3.\n"
            "5. Let's plug the numbers in:\n"
            "   - P(A) = 13/52\n"
            "   - P(B) = 12/52\n"
            "   - P(A and B) = 3/52\n"
            "   - P(A or B) = 13/52 + 12/52 - 3/52 = (13 + 12 - 3)/52 = 22/52.\n"
            "6. Simplify the fraction:\n"
            "   - 22/52 = 11/26.\n"
            "   - Decimal value: ~0.4231 (or 42.3%).\n"
            "7. Structure the explanation clearly, showing the formulas, steps, and intermediate values using LaTeX mathematical notations."
        )
        body = (
            "To find the probability that a randomly drawn card from a standard deck is either a **spade** or a **face card**, we can use the addition rule of probability.\n\n"
            "### 1. Identify the Sample Spaces\n"
            "- **Total cards in a deck:** \\(N(S) = 52\\)\n"
            "- **Spades in a deck (Event \\(A\\)):** There are 13 spades. Hence, \\(N(A) = 13\\).\n"
            "- **Face cards in a deck (Event \\(B\\)):** There are 3 face cards (Jack, Queen, King) per suit, across 4 suits. Hence, \\(N(B) = 3 \\times 4 = 12\\).\n\n"
            "### 2. Find the Intersection (Spade Face Cards)\n"
            "Some cards belong to both sets: Jack of Spades, Queen of Spades, and King of Spades. \n"
            "Let this intersection be Event \\(A \\cap B\\):\n"
            "\\[N(A \\cap B) = 3\\]\n\n"
            "### 3. Apply the Addition Rule\n"
            "The probability of the union of two events is given by:\n"
            "\\[P(A \\cup B) = P(A) + P(B) - P(A \\cap B)\\]\n\n"
            "Substitute the values:\n"
            "\\[P(A \\cup B) = \\frac{13}{52} + \\frac{12}{52} - \\frac{3}{52}\\]\n"
            "\\[P(A \\cup B) = \\frac{13 + 12 - 3}{52} = \\frac{22}{52}\\]\n\n"
            "### 4. Simplify the Result\n"
            "Reducing \\(\\frac{22}{52}\\) by dividing the numerator and denominator by 2:\n"
            "\\[P(A \\cup B) = \\frac{11}{26} \\approx 0.4231 \\text{ (or } 42.31\\%\\text{)}\\]\n\n"
            "**Conclusion:** The probability of drawing a spade or a face card is **\\(\\frac{11}{26}\\)**, which is approximately **42.3%**."
        )
    elif "box" in message_lower or "fruit" in message_lower or "label" in message_lower or "logic" in message_lower:
        thought = (
            "1. Three boxes: Box A (labeled Apples), Box B (labeled Oranges), Box C (labeled Mixed).\n"
            "2. Fact: *Every single label is incorrect*.\n"
            "   - Box labeled Apples has Oranges or Mixed.\n"
            "   - Box labeled Oranges has Apples or Mixed.\n"
            "   - Box labeled Mixed has Apples or Oranges.\n"
            "3. Let's draw a fruit from the \"Mixed\" box. Why?\n"
            "   - Since the label \"Mixed\" is wrong, it must contain either 100% Apples or 100% Oranges.\n"
            "   - If we draw a fruit and it's an Apple, then the \"Mixed\" box is 100% Apples.\n"
            "   - Now we have the remaining boxes labeled \"Apples\" and \"Oranges\".\n"
            "   - The remaining contents are \"Oranges\" and \"Mixed\".\n"
            "   - The box labeled \"Oranges\" cannot contain Oranges (as all labels are wrong). So it must contain \"Mixed\".\n"
            "   - Consequently, the box labeled \"Apples\" must contain \"Oranges\".\n"
            "4. What if we drew an Orange? By symmetry:\n"
            "   - The \"Mixed\" box is 100% Oranges.\n"
            "   - The box labeled \"Apples\" cannot contain Apples, so it must contain \"Mixed\".\n"
            "   - The box labeled \"Oranges\" must contain \"Apples\".\n"
            "5. The proof is sound and deterministic. Explain case by case with a markdown table."
        )
        body = (
            "This is a classic logical puzzle. The key to solving it is to **draw a fruit from the box labeled \"Mixed\"**.\n\n"
            "Here is the step-by-step reasoning:\n\n"
            "### 1. Understand the Rules\n"
            "- There are three boxes containing: **Apples**, **Oranges**, or **Mixed (both)**.\n"
            "- **Crucial Fact:** *All* labels are currently incorrect. This means:\n"
            "  - The box labeled **\"Apples\"** does *not* contain only apples (it contains either Oranges or Mixed).\n"
            "  - The box labeled **\"Oranges\"** does *not* contain only oranges (it contains either Apples or Mixed).\n"
            "  - The box labeled **\"Mixed\"** does *not* contain both (it contains either 100% Apples or 100% Oranges).\n\n"
            "### 2. Choose the \"Mixed\" Box First\n"
            "Since you know the box labeled \"Mixed\" is *definitely* not mixed, it must be either **100% Apples** or **100% Oranges**.\n\n"
            "- Draw one fruit from the box labeled **\"Mixed\"**.\n\n"
            "#### Case A: You draw an Apple\n"
            "1. Since this box cannot be mixed, and it contains an apple, it **must contain only Apples**. (Label this box **Apples**).\n"
            "2. You have two boxes left, labeled **\"Apples\"** and **\"Oranges\"**, and two contents left to assign: **Oranges** and **Mixed**.\n"
            "3. Look at the box labeled **\"Oranges\"**. Because all labels are wrong, this box *cannot* contain Oranges. Therefore, it **must contain Mixed**.\n"
            "4. This leaves the box labeled **\"Apples\"** to **contain only Oranges**.\n\n"
            "#### Case B: You draw an Orange\n"
            "1. Since this box cannot be mixed, and it contains an orange, it **must contain only Oranges**. (Label this box **Oranges**).\n"
            "2. You have two boxes left, labeled **\"Apples\"** and **\"Oranges\"**, and two contents left to assign: **Apples** and **Mixed**.\n"
            "3. Look at the box labeled **\"Apples\"**. Because all labels are wrong, this box *cannot* contain Apples. Therefore, it **must contain Mixed**.\n"
            "4. This leaves the box labeled **\"Oranges\"** to **contain only Apples**.\n\n"
            "### Summary Table\n\n"
            "| Label on Box | Drawn Fruit | Actual Contents | Box 2 Actual | Box 3 Actual |\n"
            "| :--- | :--- | :--- | :--- | :--- |\n"
            "| **\"Mixed\"** | Apple 🍎 | **Apples** | **\"Oranges\"** label $\\rightarrow$ **Mixed** | **\"Apples\"** label $\\rightarrow$ **Oranges** |\n"
            "| **\"Mixed\"** | Orange 🍊 | **Oranges** | **\"Apples\"** label $\\rightarrow$ **Mixed** | **\"Oranges\"** label $\\rightarrow$ **Apples** |\n\n"
            "By drawing just **one fruit from the \"Mixed\" box**, you can confidently relabel all three boxes!"
        )
    else:
        thought = (
            "1. The user prompt is general: '" + message + "'.\n"
            "2. Let's formulate a structured response that explains who I am and how I can help.\n"
            "3. Demonstrate reasoning by showing my system stats, architecture (Qwen2.5-Coder backbone), and optimal use cases (coding, math, logic).\n"
            "4. Conclude with a helpful, inviting sign-off."
        )
        body = (
            "Hello! I am Qwopus, a small language model optimized for deep, verifiable reasoning in math, coding, and STEM.\n\n"
            "To answer your request, here is a summary of how I operate:\n\n"
            "### 1. Model Capabilities\n"
            "- **Architecture:** Finetuned on top of Qwen3.5-27B using a curriculum-based SFT pipeline and Reinforcement Learning (MGPO).\n"
            "- **Reasoning Process:** I generate intermediate thoughts inside `<think>...</think>` tags to verify my hypotheses and correct mistakes before showing the final result.\n\n"
            "### 2. Suggested Prompts\n"
            "- **Math:** Ask me complex algebra, probability, or number theory questions.\n"
            "- **Coding:** Give me algorithmic coding challenges, code optimization tasks, or debugging requests.\n"
            "- **Logic:** Present me with riddles, brain teasers, or rule-based scheduling problems.\n\n"
            "Feel free to write a mathematical problem or code request to see me think through it step-by-step!"
        )
    return thought, body


def simulate_inference(message: str):
    """Simulates token-by-token streaming of the reasoning and body response."""
    thought, body = get_mock_response(message)
    full_text = f"<think>\n{thought}\n</think>\n{body}"
    
    accumulated = ""
    words = []
    current_word = ""
    for char in full_text:
        if char in (" ", "\n"):
            if current_word:
                words.append(current_word)
                current_word = ""
            words.append(char)
        else:
            current_word += char
    if current_word:
        words.append(current_word)
        
    for word in words:
        accumulated += word
        yield accumulated
        # Simulating slightly slower output for thinking steps to make it feel deliberate
        if "<think>" in accumulated and "</think>" not in accumulated:
            time.sleep(0.012 if word not in ("\n", " ") else 0.004)
        else:
            time.sleep(0.006 if word not in ("\n", " ") else 0.002)


def infer_real(message: str, system_prompt: str, temperature: float, top_p: float, max_tokens: int):
    """Executes actual model generation using the transformers library."""
    from transformers import TextIteratorStreamer, GenerationConfig
    from threading import Thread
    
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": message})
    
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    # Configure generation
    gen_kwargs = {
        "max_new_tokens": max_tokens,
        "do_sample": True if temperature > 0.0 else False,
        "top_k": None,
    }
    if temperature > 0.0:
        gen_kwargs["temperature"] = temperature
        gen_kwargs["top_p"] = top_p
        
    generation_config = GenerationConfig(**gen_kwargs)
    
    generation_kwargs = dict(
        **model_inputs,
        streamer=streamer,
        generation_config=generation_config
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    accumulated = ""
    for new_text in streamer:
        accumulated += new_text
        yield accumulated


# Define API Endpoint
# Using the wrapper helper to handle spaces.GPU decorator safely
if has_spaces:
    @app.api()
    @spaces.GPU
    def predict(
        message: str,
        system_prompt: str = "You were made by Arush",
        temperature: float = 1.0,
        top_p: float = 0.95,
        max_tokens: int = 4096
    ) -> Generator[str, None, None]:
        if is_fallback:
            yield from simulate_inference(message)
        else:
            yield from infer_real(message, system_prompt, temperature, top_p, max_tokens)
else:
    @app.api()
    def predict(
        message: str,
        system_prompt: str = "You were made by Arush",
        temperature: float = 1.0,
        top_p: float = 0.95,
        max_tokens: int = 4096
    ) -> Generator[str, None, None]:
        if is_fallback:
            yield from simulate_inference(message)
        else:
            yield from infer_real(message, system_prompt, temperature, top_p, max_tokens)


# Standard FastAPI Route to serve the frontend html
@app.get("/", response_class=HTMLResponse)
def homepage():
    try:
        with open(HTML_PATH, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read())
    except FileNotFoundError:
        return HTMLResponse(content="<h1>index.html not found! Please create the frontend page.</h1>", status_code=404)


if __name__ == "__main__":
    app.launch(show_error=True)