import os
import sys
import io
import re
import time
import json
import random
import builtins
import subprocess
import concurrent.futures

import google.generativeai as genai
from datasets import load_dataset
from google.api_core import exceptions as google_exceptions

from dotenv import load_dotenv

load_dotenv()

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY_PS")

if not GEMINI_API_KEY:
    print("❌ ERROR: GEMINI_API_KEY_PS is missing!")
else:
    GEMINI_API_KEY = GEMINI_API_KEY.strip().strip('"').strip("'")
    genai.configure(api_key=GEMINI_API_KEY)

model = genai.GenerativeModel('gemini-2.5-flash')

LEVEL_MAP = {
    "Fresh":  "A",
    "Junior": "B",
    "Senior": "C"
}

# ---------------------------------------------------------
# 2. HELPER: SAFE API CALL WITH RATE LIMIT HANDLING
# ---------------------------------------------------------

def generate_content_safe(prompt):
    retries = 0
    max_retries = 7
    while retries < max_retries:
        try:
            return model.generate_content(prompt)
        except google_exceptions.TooManyRequests:
            wait = min(60 * (2 ** retries), 300)
            print(f"⚠️ Rate Limit Hit. Cooling down for {wait}s... ({retries+1}/{max_retries})")
            time.sleep(wait)
            retries += 1
        except Exception as e:
            err = str(e)
            if "429" in err or "resource_exhausted" in err.lower():
                wait = min(60 * (2 ** retries), 300)
                print(f"⚠️ Rate Limit (429). Cooling down for {wait}s... ({retries+1}/{max_retries})")
                time.sleep(wait)
                retries += 1
            elif "503" in err or "service unavailable" in err.lower():
                wait = (2 ** retries) * 5
                print(f"⚠️ Server Busy (503). Waiting {wait}s... ({retries+1}/{max_retries})")
                time.sleep(wait)
                retries += 1
            else:
                print(f"❌ API Error: {e}")
                return None
    return None

# ---------------------------------------------------------
# 3. TIMEOUT RUNNER
# ---------------------------------------------------------

def run_func_with_timeout(func, *args, timeout_sec=2.0):
    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(func, *args)
        try:
            return future.result(timeout=timeout_sec)
        except concurrent.futures.TimeoutError:
            raise TimeoutError("Execution timed out.")

# ---------------------------------------------------------
# 4. ORACLE RUNNER (runs original dataset solution as ground truth)
# ---------------------------------------------------------

def run_oracle_subprocess(code_str, input_str):
    """Fallback: run original solution as real subprocess."""
    try:
        proc = subprocess.run(
            [sys.executable, "-c", code_str],
            input=input_str,
            capture_output=True,
            text=True,
            timeout=5.0
        )
        if proc.returncode != 0:
            return f"ORACLE_ERROR: {proc.stderr.strip()}"
        return proc.stdout.strip()
    except subprocess.TimeoutExpired:
        return "ORACLE_ERROR: subprocess timeout"
    except Exception as e:
        return f"ORACLE_ERROR: {e}"


def create_oracle_runner(code_str):
    """
    Creates a callable that runs the original dataset solution
    against any input string and returns its output.
    """
    if not code_str or not str(code_str).strip():
        return None, "No original solution provided."

    try:
        compiled_code = compile(str(code_str), '<oracle>', 'exec')
    except Exception as e:
        return None, f"Failed to compile oracle: {e}"

    def run_oracle(input_str):
        input_buffer  = io.StringIO(input_str)
        output_buffer = io.StringIO()

        class FakeBuffer:
            def read(self):
                return input_str.encode('utf-8')
            def readline(self):
                return input_buffer.readline().encode('utf-8')
            def fileno(self):
                raise io.UnsupportedOperation("fileno")

        def mocked_input(prompt=""):
            line = input_buffer.readline()
            if not line:
                raise EOFError("Oracle read more input than provided.")
            return line.strip('\n')

        old_stdin  = sys.stdin
        old_stdout = sys.stdout
        sys.stdin        = input_buffer
        sys.stdin.buffer = FakeBuffer()
        sys.stdout       = output_buffer

        namespace = {
            "__builtins__": builtins,
            "__name__":     "__main__",
            "sys":          sys,
            "io":           io,
            "math":         __import__('math'),
            "input":        mocked_input,
            "raw_input":    mocked_input,
            "exit":         lambda *a: (_ for _ in ()).throw(SystemExit(0)),
            "quit":         lambda *a: (_ for _ in ()).throw(SystemExit(0)),
        }

        try:
            exec(compiled_code, namespace)
            return output_buffer.getvalue().strip()
        except SystemExit:
            return output_buffer.getvalue().strip()
        except io.UnsupportedOperation:
            print("   ⚠️  Oracle needs real stdin — switching to subprocess fallback.")
            return run_oracle_subprocess(str(code_str), input_str)
        except EOFError as e:
            return f"ORACLE_ERROR: {e}"
        except Exception as e:
            return f"ORACLE_ERROR: {e}"
        finally:
            sys.stdin  = old_stdin
            sys.stdout = old_stdout

    return run_oracle, "Success"

# ---------------------------------------------------------
# 5. FETCH SEED PROBLEM FROM DEEPMIND/CODE_CONTESTS
def get_random_seed_problem(target_difficulty_char):
    print(f"🔍 Searching for a Level '{target_difficulty_char}' seed from DeepMind/code_contests...")
    try:
        # Load the stream INSIDE the function so it doesn't block server startup
        dataset = load_dataset("deepmind/code_contests", split="train", streaming=True)
        # Small buffer (50) so it starts instantly
        shuffled = dataset.shuffle(seed=random.randint(0, 100000), buffer_size=50)
        iterator = iter(shuffled)

        attempts     = 0
        max_attempts = 2000

        while attempts < max_attempts:
            try:
                p = next(iterator)
                attempts += 1

                # ── Filter 1: Codeforces source only
                if p.get('source') != 2:
                    continue

                # ── Filter 2: Correct difficulty index
                if p.get('cf_index', '') != target_difficulty_char:
                    continue

                # ── Filter 3: Description long enough
                desc = p.get('description', '')
                if len(desc) < 300:
                    continue

                # ── Filter 4: Skip image-based problems
                if '<image>' in desc.lower():
                    continue

                # ── Filter 5: Must have at least one Python solution (FIXED)
                solutions = p.get('solutions', {})
                original_solution = None
                
                # Codeforces IDs for Python: 3 (Py3), 2 (Py2), 31 (Py3.4), 41 (PyPy3), 70 (PyPy3.7)
                python_ids = {2, 3, 31, 41, 70}

                if isinstance(solutions, dict) and 'language' in solutions and 'solution' in solutions:
                    for lang_id, sol_code in zip(solutions['language'], solutions['solution']):
                        if lang_id in python_ids:  
                            original_solution = sol_code
                            break
                            
                elif isinstance(solutions, list):
                    for sol in solutions:
                        if isinstance(sol, dict) and sol.get('language') in python_ids:
                            original_solution = sol.get('solution', '')
                            break

                if not original_solution:
                    continue

                print(f"✅ Found Seed: {p['name']} (Difficulty: {target_difficulty_char}) after {attempts} attempts")
                return {
                    "name":        p['name'],
                    "description": desc,
                    "solution":    original_solution,
                }

            except StopIteration:
                break
            except Exception:
                continue

    except Exception as e:
        print(f"❌ Dataset Error: {e}")

    print(f"❌ Could not find a valid seed.")
    return None
# ---------------------------------------------------------
# 6. REWRITE PROBLEM STORY WITH GEMINI
# ---------------------------------------------------------

def generate_problem_data(seed):
    """
    Sends the original problem + solution to Gemini.
    Gemini rewrites the story (theme, names, context) but keeps the
    exact same algorithmic logic and I/O format.
    """
    print(f"✍️  Rewriting story for: '{seed['name']}'...")

    prompt = f"""
    You are a Competitive Programming Problem Setter.

    Below is an original problem and its verified correct solution.
    Your job is to create a BRAND NEW problem by changing ONLY the story/theme —
    keep the exact same algorithmic logic, input/output format, and constraints.

    === ORIGINAL PROBLEM ===
    {seed['description'][:2000]}

    === ORIGINAL SOLUTION (keep the logic identical) ===
    ```python
    {seed['solution'][:1500]}
    ```

    INSTRUCTIONS:
    1. Change the background story completely (e.g. space, magic, robots, cooking).
    2. Keep every constraint, formula, and edge case exactly the same.
    3. Keep the input/output format exactly the same (same number of lines, same types).
    4. Do NOT reveal the algorithm name in the problem text.

    Output STRICTLY in this format (no extra text outside tags):
    [TITLE_START]...[TITLE_END]
    [DESC_START]...[DESC_END]
    [INPUT_START]...[INPUT_END]
    [OUTPUT_START]...[OUTPUT_END]
    """

    response = generate_content_safe(prompt)
    if not response:
        return None

    text = response.text
    try:
        return {
            "title":       text.split("[TITLE_START]")[1].split("[TITLE_END]")[0].strip(),
            "description": text.split("[DESC_START]")[1].split("[DESC_END]")[0].strip(),
            "input_fmt":   text.split("[INPUT_START]")[1].split("[INPUT_END]")[0].strip(),
            "output_fmt":  text.split("[OUTPUT_START]")[1].split("[OUTPUT_END]")[0].strip(),
        }
    except Exception:
        print("⚠️  Failed to parse rewritten problem format.")
        return None

# ---------------------------------------------------------
# 7. GENERATE CODE COMPONENTS (with oracle hint + error feedback)
# ---------------------------------------------------------

def generate_code_components(problem_data, original_solution="", previous_error=""):
    """
    Asks Gemini to generate:
      solve_optimized, solve_brute,
      generate_small_test, generate_large_test, generate_edge_test,
      validate_output
    Using the original solution as the logic source of truth.
    """
    print("⚙️  Generating Logic Components (Python)...")

    oracle_hint = ""
    if original_solution and original_solution.strip():
        oracle_hint = f"""
=== REFERENCE SOLUTION (YOUR LOGIC SOURCE OF TRUTH) ===
The following is the verified correct solution from the dataset.
Your `solve_optimized` AND `solve_brute` MUST produce EXACTLY the same
output as this code for every possible input.
Do NOT invent new logic. Only adapt the I/O to accept a plain string argument.
```python
{original_solution}
```
========================================================
"""

    error_hint = ""
    if previous_error:
        error_hint = f"""
=== PREVIOUS ATTEMPT FAILED — FIX THIS ===
The last generated code failed with:
\"\"\"{previous_error}\"\"\"

Common causes:
- solve_optimized returned a different value than the oracle on some edge case.
- Integer division done with / instead of //.
- A string literal like "No"/"Yes"/"Impossible" was changed.
- generate_small_test or generate_edge_test produced malformed input.
- This is a constructive problem — make sure validate_output checks ALL problem
  rules, not just one specific answer.

Study the error carefully and fix the root cause.
==========================================
"""

    prompt = f"""
    Problem Title: {problem_data['title']}
    Description: {problem_data['description']}
    Input Format: {problem_data['input_fmt']}

    {oracle_hint}
    {error_hint}

    Generate a SINGLE Python code block with these 6 functions:

    1. `def solve_optimized(input_str):`
       - Efficient O(N) or O(N log N) solution.
       - Read from `input_str` argument (NOT stdin).

    2. `def solve_brute(input_str):`
       - Simple correct brute-force for verification.
       - Read from `input_str` argument (NOT stdin).

    === CRITICAL RULES FOR TEST GENERATORS ===
    - DATA INTEGRITY: If format requires N then N integers, generate EXACTLY N integers.
    - TYPE STRICTNESS: If format requires integers, do not generate floats or letters.
    - FORMAT EXACTNESS: Use \\n to separate lines exactly as the Input Format specifies.
      Do NOT add trailing newlines or extra spaces at the end.
    ==========================================

    3. `def generate_small_test():`
       - Returns a RANDOM input string based on the format. N = 1 to 10.
       - MUST use `random` so every call produces a DIFFERENT output.

    4. `def generate_large_test():`
       - Returns a RANDOM input string. N = 100 to 1000.
       - MUST use `random` module to ensure uniqueness.

    5. `def generate_edge_test():`
       - Returns a RANDOM edge case input string.
       - Examples: N=1, all identical elements, minimum possible values.
       - Keep N VERY SMALL (N < 20) to prevent brute-force from freezing.
       - Use `random.choice` to pick between different edge case types.

    6. `def validate_output(input_str, output_str):`
       - Many problems have MULTIPLE valid answers (constructive problems).
       - Check whether output_str is ANY valid answer per the problem rules.
       - Parse input_str and output_str, verify ALL constraints.
       - Return True if valid, False otherwise.
       - For single-answer problems a basic sanity check is fine.

    IMPORTANT:
    - Import `random` inside the code block.
    - Do NOT use external libraries (like numpy).
    - Wrap the entire code in ONE ```python``` block.
    """

    response = generate_content_safe(prompt)
    if not response:
        return None

    text  = response.text
    match = re.search(r"```python(.*?)```", text, re.DOTALL)
    return match.group(1).strip() if match else text.strip()

# ---------------------------------------------------------
# 8. VERIFY & JUDGE (oracle-based)
# ---------------------------------------------------------

def verify_and_judge(code_block, original_solution=""):
    """
    Verifies the generated logic against the original solution (oracle).
    Returns (is_good, status_message, namespace, hidden_tests).
    """
    print("🔬 Running Verification Loop (Oracle + Brute Force)...")

    oracle_runner, oracle_status = create_oracle_runner(original_solution)
    if not oracle_runner:
        print(f"   ⚠️  No oracle available ({oracle_status}). Falling back to brute-only check.")

    namespace = {}
    try:
        exec(code_block, namespace)
        solve_opt   = namespace.get('solve_optimized')
        solve_brute = namespace.get('solve_brute')
        gen_small   = namespace.get('generate_small_test')
        gen_large   = namespace.get('generate_large_test')
        gen_edge    = namespace.get('generate_edge_test')
        validate_fn = namespace.get('validate_output')

        if not all([solve_opt, solve_brute, gen_small, gen_large, gen_edge]):
            return False, "Missing required functions in generated code.", None, []

        hidden_tests = []
        seen_inputs  = set()

        # ── Token-level comparison (handles float answers)
        def compare_tokens(a, b):
            ta, tb = a.split(), b.split()
            if len(ta) != len(tb):
                return False
            for x, y in zip(ta, tb):
                try:
                    if abs(float(x) - float(y)) > 1e-6:
                        return False
                except ValueError:
                    if x != y:
                        return False
            return True

        def is_correct(inp, candidate, oracle_out):
            """True if candidate is a valid answer for inp."""
            if compare_tokens(candidate, oracle_out):
                return True
            if validate_fn:
                try:
                    return bool(run_func_with_timeout(validate_fn, inp, candidate, timeout_sec=2.0))
                except Exception:
                    return False
            return False

        def test_one(inp, check_brute=True):
            try:
                out_opt = str(run_func_with_timeout(solve_opt, inp, timeout_sec=2.0)).strip()
            except TimeoutError:
                return False, "TIMEOUT: solve_optimized", ""

            # ── Compare against oracle if available
            if oracle_runner:
                try:
                    out_oracle = str(run_func_with_timeout(oracle_runner, inp, timeout_sec=5.0)).strip()
                except TimeoutError:
                    return False, "TIMEOUT: oracle", ""

                if "ORACLE_ERROR" in out_oracle:
                    return False, f"Oracle crashed: {out_oracle}", ""

                if not is_correct(inp, out_opt, out_oracle):
                    preview = inp[:300].replace('\n', '\\n')
                    return False, (
                        f"Mismatch vs Oracle!\n"
                        f"  Input    : {preview}\n"
                        f"  Oracle   : {out_oracle[:120]}\n"
                        f"  Optimized: {out_opt[:120]}"
                    ), out_oracle   # return oracle output as expected

                expected = out_oracle  # oracle is ground truth

            else:
                expected = out_opt    # no oracle — use optimized as reference

            # ── Also check brute force matches
            if check_brute:
                try:
                    out_brute = str(run_func_with_timeout(solve_brute, inp, timeout_sec=4.0)).strip()
                except TimeoutError:
                    return False, "TIMEOUT: solve_brute", ""

                if not is_correct(inp, out_brute, expected):
                    preview = inp[:300].replace('\n', '\\n')
                    return False, (
                        f"Brute Mismatch!\n"
                        f"  Input    : {preview}\n"
                        f"  Expected : {expected[:120]}\n"
                        f"  Brute    : {out_brute[:120]}"
                    ), ""

            return True, "Passed", expected

        # ── Test generation plan: (generator_fn, count, check_brute?)
        print("   Generating tests (Small x5, Large x5, Edge x3)...")
        plan = [
            (gen_small, 5,  True),
            (gen_large, 5,  False),
            (gen_edge,  3,  True),
        ]

        for gen_fn, target_count, check_b in plan:
            collected = 0
            attempts  = 0
            while collected < target_count and attempts < 40:
                attempts += 1
                try:
                    inp = str(run_func_with_timeout(gen_fn, timeout_sec=2.0)).strip()
                    if not inp or inp in seen_inputs:
                        continue

                    passed, msg, expected = test_one(inp, check_brute=check_b)
                    if not passed:
                        return False, msg, None, []

                    seen_inputs.add(inp)
                    hidden_tests.append({
                        "id":              len(hidden_tests),
                        "input":           inp,
                        "expected_output": expected
                    })
                    collected += 1
                except Exception:
                    continue

        print(f"   ✅ Total Tests Verified: {len(hidden_tests)}")
        return True, "Verified", namespace, hidden_tests

    except Exception as e:
        return False, f"Exec Error: {e}", None, []

# ---------------------------------------------------------
# 9. GENERATE SAMPLE TESTS (visible to user)
# ---------------------------------------------------------

def create_sample_tests(namespace):
    """Creates 3 visible sample tests from the small test generator."""
    samples  = []
    seen     = set()
    gen_fn   = namespace.get('generate_small_test')
    solve_fn = namespace.get('solve_optimized')

    if not gen_fn or not solve_fn:
        return samples

    print("   Generating 3 sample tests...")
    tries = 0
    while len(samples) < 3 and tries < 25:
        tries += 1
        try:
            inp = str(gen_fn()).strip()
            if inp in seen:
                continue
            out = str(solve_fn(inp)).strip()
            seen.add(inp)
            samples.append({"input": inp, "output": out})
        except Exception:
            continue

    return samples

# ---------------------------------------------------------
# 10. GENERATE TARGET SOLUTIONS (Python, C++, JS)
# ---------------------------------------------------------

def generate_target_solution(problem_data, verified_python_code):
    """
    Generates clean, runnable solutions for Python 3, C++, and JavaScript
    from the verified internal logic.
    """
    print("🌐 Generating Target Solutions (Python / C++ / JS)...")

    prompt = f"""
    I have a competitive programming problem and its verified internal logic.
    Write a clean, optimal, STANDALONE solution in Python 3, C++, and JavaScript (Node.js).

    PROBLEM DESCRIPTION:
    {problem_data['description']}

    INPUT FORMAT:  {problem_data['input_fmt']}
    OUTPUT FORMAT: {problem_data['output_fmt']}

    INTERNAL LOGIC (SOURCE OF TRUTH — translate this exactly):
    {verified_python_code}

    CRITICAL TRANSLATION RULES:
    1. EXACT LOGIC: 1:1 translation of `solve_optimized`. Do NOT invent new logic.
    2. EXACT OUTPUT: Look at string literals in `solve_optimized` (e.g. "No collision",
       "Impossible"). ALL languages MUST print that EXACT string.
    3. EXACT FORMAT: Space-separated numbers stay space-separated. No [] or commas.
    4. OVERFLOW: C++ MUST use `long long`. JS MUST use BigInt if numbers exceed 2^53.
    5. JAVASCRIPT RULES (CRITICAL):
       - Integer division: `Math.floor(A / B)` — NEVER plain `A / B`.
       - Sorting:          `.sort((a, b) => a - b)` — NEVER empty `.sort()`.
       - Parsing:          `const tokens = input.trim().split(/\\s+/).filter(Boolean);`
    6. INPUT READING:
       - JavaScript: `fs.readFileSync(0, 'utf-8')`.
       - Python:     `sys.stdin.read()`.
    7. CLEANUP: Only the final runnable solution. No test generators, no `import random`.

    OUTPUT STRICTLY IN THIS FORMAT (@@@ as separator):
    @@@PYTHON@@@
    <python code here>
    @@@CPP@@@
    <cpp code here>
    @@@JS@@@
    <javascript code here>
    """

    response  = generate_content_safe(prompt)
    solutions = {
        "python": verified_python_code,
        "cpp":    "// C++ solution failed to generate.",
        "js":     "// JavaScript solution failed to generate."
    }

    if not response:
        return solutions

    text = response.text

    def clean(raw):
        raw = re.sub(r"```[a-zA-Z+]*\n?", "", raw)
        return re.sub(r"```", "", raw).strip()

    py_m  = re.search(r"@@@\s*PYTHON\s*@@@(.*?)(?=@@@\s*(?:CPP|JS)\s*@@@|$)",  text, re.DOTALL | re.IGNORECASE)
    cpp_m = re.search(r"@@@\s*CPP\s*@@@(.*?)(?=@@@\s*(?:JS|PYTHON)\s*@@@|$)",  text, re.DOTALL | re.IGNORECASE)
    js_m  = re.search(r"@@@\s*JS\s*@@@(.*?)(?=@@@\s*(?:PYTHON|CPP)\s*@@@|$)",  text, re.DOTALL | re.IGNORECASE)

    if py_m:  solutions["python"] = clean(py_m.group(1))
    if cpp_m: solutions["cpp"]    = clean(cpp_m.group(1))
    if js_m:  solutions["js"]     = clean(js_m.group(1))

    if not any([py_m, cpp_m, js_m]):
        print("⚠️  generate_target_solution: No @@@ markers found in response.")

    return solutions

# ---------------------------------------------------------
# 11. API ENTRY POINT
# ---------------------------------------------------------

def generate_problem_for_api(level: str, lang: str) -> dict | None:
    """
    Main function called by FastAPI.
    Returns a complete problem dictionary or None on failure.
    """
    target_char = LEVEL_MAP.get(level, "A")

    # ── Step 1: Fetch seed from DeepMind dataset
    seed = get_random_seed_problem(target_char)
    if not seed:
        print("❌ Could not fetch a valid seed problem.")
        return None

    # ── Step 2: Rewrite story with Gemini
    new_prob = generate_problem_data(seed)
    if not new_prob:
        print("❌ Story rewrite failed.")
        return None

    # ── Step 3: Generate & verify logic with retry loop
    MAX_ATTEMPTS = 3
    last_error   = ""
    code_block   = None
    namespace    = None
    hidden_tests = []
    is_valid     = False

    for attempt in range(MAX_ATTEMPTS):
        print(f"\n🔄 Code generation attempt {attempt + 1}/{MAX_ATTEMPTS}...")

        code_block = generate_code_components(
            new_prob,
            original_solution=seed['solution'],
            previous_error=last_error
        )

        if not code_block:
            last_error = "generate_code_components returned None (API error)."
            print(f"   ⚠️  Empty code block — waiting 15s before retry...")
            time.sleep(15)
            continue

        is_good, status, namespace, hidden_tests = verify_and_judge(
            code_block, original_solution=seed['solution']
        )

        if is_good:
            is_valid = True
            print(f"   ✅ Verified on attempt {attempt + 1}.")
            break

        last_error = status
        print(f"   ❌ Attempt {attempt + 1} failed: {status[:200]}")
        if attempt < MAX_ATTEMPTS - 1:
            time.sleep(5)

    if not is_valid:
        print(f"❌ All {MAX_ATTEMPTS} attempts failed. Aborting.")
        return None

    # ── Step 4: Generate target solutions
    target_sols = generate_target_solution(new_prob, code_block)

    # ── Step 5: Build and return result
    return {
        "title":         new_prob['title'],
        "description":   new_prob['description'],
        "input_format":  new_prob['input_fmt'],
        "output_format": new_prob['output_fmt'],
        "samples":       create_sample_tests(namespace),
        "test_cases":    hidden_tests,
        "solution_code": target_sols,   # dict: {python, cpp, js}
    }