import torch
import torch.nn.functional as F
import json
import sys
import math
import ast
import os
import time
import subprocess
import tempfile
from pathlib import Path


class CausalSelfAttention(torch.nn.Module):
    def __init__(self, d_model, n_heads, dropout, context_length):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = torch.nn.Linear(d_model, 3 * d_model)
        self.proj = torch.nn.Linear(d_model, d_model)
        self.attn_dropout = torch.nn.Dropout(dropout)
        self.resid_dropout = torch.nn.Dropout(dropout)
        self.register_buffer("mask", torch.tril(torch.ones(context_length, context_length)).unsqueeze(0).unsqueeze(0))

    def forward(self, x):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        attn = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_dim))
        attn = attn.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        attn = self.attn_dropout(attn)
        out = attn @ v
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        out = self.proj(out)
        out = self.resid_dropout(out)
        return out


class MLP(torch.nn.Module):
    def __init__(self, d_model, d_ff, dropout):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(d_model, d_ff),
            torch.nn.GELU(),
            torch.nn.Linear(d_ff, d_model),
            torch.nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class TransformerBlock(torch.nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout, context_length):
        super().__init__()
        self.ln1 = torch.nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, dropout, context_length)
        self.ln2 = torch.nn.LayerNorm(d_model)
        self.mlp = MLP(d_model, d_ff, dropout)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


def format_instruction(instruction, extra_input=""):
    instruction = (instruction or "").strip()
    extra_input = (extra_input or "").strip()
    if extra_input and extra_input.lower() != "not applicable":
        return f"### Instruction:\n{instruction}\n\n### Input:\n{extra_input}\n\n### Response:\n"
    return f"### Instruction:\n{instruction}\n\n### Response:\n"


_FOREIGN_MARKERS = (
    "#include", "void main", "int main(", "public static void",
    "System.out.println", "console.log", "function ", "</",
    "<?php", "using namespace", "fmt.Println", "package main",
    "fn main", "<html", "<script", "CREATE TABLE", "SELECT ", "=>",
)


def extract_code(text):
    text = (text or "").strip()
    if "```" not in text:
        return text.strip()

    def _drop_lang_label(block):
        lines = block.split("\n")
        if lines and lines[0].strip() and len(lines[0].strip()) <= 12 \
                and not any(ch in lines[0] for ch in " \t=()[]{}:;"):
            lines = lines[1:]
        return "\n".join(lines).strip("\n")

    parts = text.split("```")
    blocks = []
    for i in range(1, len(parts), 2):
        blocks.append(_drop_lang_label(parts[i]))
    if blocks:
        return "\n\n".join(b.strip("\n") for b in blocks).strip()

    return _drop_lang_label(parts[1]).strip()


def looks_like_python(code):
    head = (code or "")[:3000].lower()
    return not any(m.lower() in head for m in _FOREIGN_MARKERS)


def check_syntax(code):
    if not (code or "").strip():
        return False, "model returned no code (empty response)"

    try:
        ast.parse(code)
        return True, None
    except (SyntaxError, ValueError) as e:
        if not looks_like_python(code):
            return False, ("this doesn't look like Python code — syntax checking and "
                           "execution are only supported for Python")
        if isinstance(e, ValueError):
            return False, f"failed to parse code: {e}"

        lines = (code or "").splitlines()
        lineno = e.lineno or 1
        offset = e.offset or 1
        out = [f"SyntaxError: {e.msg} (line {lineno}, column {offset})"]
        if 1 <= lineno <= len(lines):
            bad_line = lines[lineno - 1]
            caret_pos = min(max(offset, 1), len(bad_line) + 1) - 1
            out.append(f"    {lineno:>4} | {bad_line}")
            out.append(f"         | {' ' * caret_pos}^")
            if lineno >= len(lines):
                out.append("    (looks like the code was cut off by the generation limit — "
                           "try increasing code_max_new_tokens)")
        return False, "\n".join(out)


def run_python_code(code, timeout=10.0):
    fd, path = tempfile.mkstemp(suffix=".py", prefix="cortex_run_")
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            f.write(code)
        env = {**os.environ, "PYTHONIOENCODING": "utf-8"}
        proc = subprocess.run(
            [sys.executable, "-u", path],
            stdin=subprocess.DEVNULL,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=timeout,
            env=env,
        )
        return proc.returncode, (proc.stdout or "") + (proc.stderr or ""), False
    except subprocess.TimeoutExpired as e:
        partial = ""
        for stream in (e.stdout, e.stderr):
            if not stream:
                continue
            if isinstance(stream, bytes):
                stream = stream.decode("utf-8", "replace")
            partial += stream
        return -1, partial, True
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass


class TinyGPT(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        vocab_size = config["tokenizer_vocab_size"] + 10
        self.token_emb = torch.nn.Embedding(vocab_size, config["d_model"])
        self.pos_emb = torch.nn.Embedding(config["context_length"], config["d_model"])
        self.drop = torch.nn.Dropout(config["dropout"])
        self.blocks = torch.nn.ModuleList([
            TransformerBlock(config["d_model"], config["n_heads"], config["d_ff"], config["dropout"], config["context_length"])
            for _ in range(config["n_layers"])
        ])
        self.ln_f = torch.nn.LayerNorm(config["d_model"])
        self.head = torch.nn.Linear(config["d_model"], vocab_size, bias=False)
        self.token_emb.weight = self.head.weight

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(0, T, device=idx.device).unsqueeze(0)
        x = self.token_emb(idx) + self.pos_emb(pos)
        x = self.drop(x)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=0)
        return logits, loss


def find_model_file():
    here = Path(".")

    pt_files = list(here.glob("*.pt"))

    for name in ["best_model.pt", "final_model.pt"]:
        if name in [f.name for f in pt_files]:
            return here / name

    if pt_files:
        return pt_files[0]

    return None


def main():
    device = torch.device("cuda")

    model_path = find_model_file()
    if model_path is None:
        print("❌ No .pt model file found! Put this script in the same folder as your model.")
        sys.exit(1)

    if len(sys.argv) > 1:
        model_path = Path(sys.argv[1])

    print(f"📂 Loading model from: {model_path.name}")

    ckpt = torch.load(model_path, map_location=device, weights_only=False)

    if "config" in ckpt and "tokenizer" in ckpt:
        config = ckpt["config"]
        from tokenizers import Tokenizer
        tokenizer = Tokenizer.from_str(ckpt["tokenizer"])
        print("📦 Loaded config + tokenizer from checkpoint")
    else:
        here = model_path.parent
        config_path = here / "config.json"
        tokenizer_path = here / "tokenizer.json"

        if not config_path.exists():
            print(f"❌ config.json not found next to model!")
            sys.exit(1)
        if not tokenizer_path.exists():
            print(f"❌ tokenizer.json not found next to model!")
            sys.exit(1)

        with open(config_path) as f:
            config = json.load(f)
        from tokenizers import Tokenizer
        tokenizer = Tokenizer.from_file(str(tokenizer_path))
        print("📦 Loaded config + tokenizer from separate files")

    model = TinyGPT(config).to(device)
    model.load_state_dict(ckpt["model"])
    model.eval()

    n_params = sum(p.numel() for p in model.parameters())
    step = ckpt.get("step", "?")
    val_loss = ckpt.get("val_loss", "?")
    if isinstance(val_loss, float):
        val_loss = f"{val_loss:.4f}"

    print(f"✅ Cortex_2 loaded!")
    print(f"   Parameters: {n_params / 1e6:.1f}M")
    print(f"   Step: {step}")
    print(f"   Val loss: {val_loss}")
    print(f"   Device: {device}")

    dataset_mode = config.get("dataset_mode", "stories")
    is_chat_model = dataset_mode == "chat"
    is_code_model = dataset_mode == "code"

    if is_chat_model:
        print(f"   Mode: 💬 conversational (dataset_mode=chat)")
    elif is_code_model:
        print(f"   Mode: 🧑‍💻 code (dataset_mode=code)")
    else:
        print(f"   Mode: 📖 story completion (dataset_mode=stories)")

    print()
    print("💬 Type a prompt and press Enter. Type 'quit' to exit.")
    if is_chat_model:
        print("   (type 'reset' to clear conversation history)")
        print("   (type 'temp 0.9' to change temperature, current default: 0.8)")
    if is_code_model:
        print("   Describe a task, e.g.: 'Write a function that reverses a string'.")
        print("   (to set a separate 'Input:', type: task || input)")
        print("   (type 'temp 0.5' to change temperature, current default: 0.5)")
        print()
        print("   Code mode commands:")
        print("     run         — run the last generated code")
        print("     save        — save the last code to generated_code_NN.py")
        print("     autocheck   — auto-regenerate on syntax error")
        print("     timeout N   — code execution timeout in seconds")
        print("   After generation the code is syntax-checked, and clean code can be")
        print("   run directly from the chat (y when asked 'Run?').")
    print("=" * 50)

    bos_id = tokenizer.token_to_id("<bos>")
    eos_id = tokenizer.token_to_id("<eos>")
    context_length = config["context_length"]

    history_lines = []
    temperature = 0.5 if is_code_model else 0.8
    code_max_new_tokens = 400
    code_top_k = 40

    last_code = None
    run_timeout = 10.0
    autocheck = True
    max_auto_attempts = 3

    def generate_code(instruction, extra_input=""):
        text_prompt = format_instruction(instruction, extra_input)
        ids = tokenizer.encode(text_prompt).ids
        idx = torch.tensor([[bos_id] + ids], dtype=torch.long, device=device)
        prompt_len = idx.shape[1]

        t0 = time.time()
        n_tokens = 0
        with torch.no_grad():
            for _ in range(code_max_new_tokens):
                idx_cond = idx[:, -context_length:]
                logits, _ = model(idx_cond)
                logits = logits[:, -1, :] / temperature
                if code_top_k:
                    kth = torch.topk(logits, code_top_k).values[:, -1, None]
                    logits = logits.masked_fill(logits < kth, float("-inf"))
                probs = F.softmax(logits, dim=-1)
                next_id = torch.multinomial(probs, num_samples=1)
                idx = torch.cat([idx, next_id], dim=1)
                n_tokens += 1
                if next_id.item() == eos_id:
                    break
        print(f"   ⏳ generated {n_tokens} tokens in {time.time() - t0:.1f}s")
        return tokenizer.decode(idx[0, prompt_len:].tolist())

    def execute_code(code):
        print("─" * 50)
        print(f"▶ Running code (separate process, timeout {run_timeout:.0f}s, stdin closed)...")
        rc, output, timed_out = run_python_code(code, run_timeout)
        if timed_out:
            print(f"⏱ Timeout exceeded ({run_timeout:.0f}s) — process stopped.")
            if output.strip():
                print("📤 Output before stopping:")
                print(output.rstrip())
            print("   Hint: if the code waits for input(), it will never finish —")
            print("   interactive input is not available when running from chat.")
        elif rc == 0:
            if output.strip():
                print("📤 Program output:")
                print(output.rstrip())
            else:
                print("📤 Program finished with no output.")
            print("✅ Code ran without errors (exit code 0).")
        else:
            if output.strip():
                print("📤 Program output:")
                print(output.rstrip())
            if "EOFError" in output:
                print("   Hint: the code called input() — input is not available when running from chat.")
            print(f"❌ Program finished with an error (exit code {rc}).")
        print("─" * 50)

    # Chat loop
    while True:
        try:
            prompt = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Bye!")
            break

        if prompt.lower() == "quit":
            print("👋 Bye!")
            break
        if is_chat_model and prompt.lower() == "reset":
            history_lines = []
            print("🔄 Conversation history cleared.")
            continue
        if (is_chat_model or is_code_model) and prompt.lower().startswith("temp"):
            parts = prompt.split()
            if len(parts) == 2:
                try:
                    new_temp = float(parts[1])
                    if new_temp <= 0:
                        print("⚠️ Temperature must be greater than 0.")
                    else:
                        temperature = new_temp
                        print(f"🌡️ Temperature set to: {temperature}")
                except ValueError:
                    print("⚠️ Could not parse the number. Example: temp 0.9")
            else:
                print(f"🌡️ Current temperature: {temperature} (example to change: temp 0.9)")
            continue

        if is_code_model and prompt.lower() in ("run", "r"):
            if not last_code:
                print("⚠️ Nothing to run yet — generate some code first.")
                continue
            ok, err = check_syntax(last_code)
            if not ok:
                print(f"❌ The last code has a syntax error, cannot run it:\n{err}")
                continue
            execute_code(last_code)
            continue

        if is_code_model and prompt.lower() == "save":
            if not last_code:
                print("⚠️ Nothing to save yet — generate some code first.")
                continue
            n = 1
            while (Path.cwd() / f"generated_code_{n:02d}.py").exists():
                n += 1
            save_path = Path.cwd() / f"generated_code_{n:02d}.py"
            save_path.write_text(last_code, encoding="utf-8")
            print(f"💾 Code saved: {save_path}")
            continue

        if is_code_model and prompt.lower().startswith("autocheck"):
            parts = prompt.split()
            if len(parts) == 2 and parts[1].lower() in ("on", "off"):
                autocheck = parts[1].lower() == "on"
                state = "on" if autocheck else "off"
                print(f"🔄 Auto-regenerate on error: {state} (max attempts: {max_auto_attempts})")
            else:
                state = "on" if autocheck else "off"
                print(f"🔄 Auto-regenerate is currently: {state} (example: autocheck off)")
            continue

        if is_code_model and prompt.lower().startswith("timeout"):
            parts = prompt.split()
            if len(parts) == 2:
                try:
                    val = float(parts[1])
                    if val <= 0:
                        print("⚠️ Timeout must be greater than 0.")
                    else:
                        run_timeout = val
                        print(f"⏱ Code execution timeout: {run_timeout:.0f}s")
                except ValueError:
                    print("⚠️ Could not parse the number. Example: timeout 15")
            else:
                print(f"⏱ Current execution timeout: {run_timeout:.0f}s (example: timeout 15)")
            continue

        if not prompt:
            continue

        if is_chat_model:
            history_lines.append(f"User: {prompt}")
            history_lines.append("Bot:")
            full_text = "\n".join(history_lines)

            ids = tokenizer.encode(full_text).ids
            idx = torch.tensor([[bos_id] + ids], dtype=torch.long, device=device)

            tokens_before_gen = idx.shape[1]

            if idx.shape[1] > context_length:
                idx = idx[:, -context_length:]

            generated_ids = []
            with torch.no_grad():
                for _ in range(200):
                    idx_cond = idx[:, -context_length:]
                    logits, _ = model(idx_cond)
                    logits = logits[:, -1, :]
                    probs = F.softmax(logits / temperature, dim=-1)
                    next_id = torch.multinomial(probs, num_samples=1)
                    idx = torch.cat([idx, next_id], dim=1)
                    generated_ids.append(next_id.item())

                    if next_id.item() == eos_id:
                        break

                    partial_text = tokenizer.decode(generated_ids)
                    normalized = partial_text.replace(" :", ":").replace(" ,", ",")
                    if "User:" in normalized:
                        break

            reply_text = tokenizer.decode(generated_ids)
            normalized_reply = reply_text.replace(" :", ":")
            if "User:" in normalized_reply:
                cut_pos = normalized_reply.index("User:")
                reply_text = reply_text.split("User :")[0].split("User:")[0].strip()
            else:
                reply_text = reply_text.strip()

            print(f"Cortex_2: {reply_text}")

            history_lines[-1] = f"Bot: {reply_text}"

            tokens_used = min(tokens_before_gen + len(generated_ids), context_length)
            pct = tokens_used / context_length * 100
            print(f"📊 Context: {tokens_used}/{context_length} tokens ({pct:.1f}%)")

        elif is_code_model:
            if "||" in prompt:
                instruction, extra_input = prompt.split("||", 1)
            else:
                instruction, extra_input = prompt, ""
            instruction = instruction.strip()

            code_text = extract_code(generate_code(instruction, extra_input))
            ok, err = check_syntax(code_text)

            attempt = 1
            while not ok and autocheck and attempt < max_auto_attempts:
                attempt += 1
                print(f"🔄 Attempt {attempt}/{max_auto_attempts}: code has an error, regenerating...")
                code_text = extract_code(generate_code(instruction, extra_input))
                ok, err = check_syntax(code_text)

            print(f"Cortex_2:\n{code_text}")
            last_code = code_text

            if ok:
                print("✅ Syntax: no errors found")
                try:
                    ans = input("▶ Run this code? [y/N]: ").strip().lower()
                except (EOFError, KeyboardInterrupt):
                    ans = ""
                if ans in ("y", "yes"):
                    execute_code(code_text)
            else:
                print(f"❌ Syntax: error found!\n{err}")
                if not autocheck:
                    print("   Hint: enable autocheck on — the chat will try to")
                    print("   regenerate the code automatically on error.")

        else:
            ids = tokenizer.encode(prompt).ids
            idx = torch.tensor([[bos_id] + ids], dtype=torch.long, device=device)

            with torch.no_grad():
                for _ in range(750):
                    idx_cond = idx[:, -context_length:]
                    logits, _ = model(idx_cond)
                    logits = logits[:, -1, :]
                    probs = F.softmax(logits / 0.8, dim=-1)
                    next_id = torch.multinomial(probs, num_samples=1)
                    idx = torch.cat([idx, next_id], dim=1)
                    if next_id.item() == eos_id:
                        break

            text = tokenizer.decode(idx[0].tolist())
            print(f"Cortex_2: {text}")


if __name__ == "__main__":
    main()