src/pageparse/dsa_visualizer.py

from __future__ import annotations

import json
import re
import urllib.request
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any

from pageparse.config import settings


@dataclass
class VisualizationNode:
    id: str
    label: str
    value: str | None = None
    color: str = "#4a90d9"
    x: float = 0
    y: float = 0


@dataclass
class VisualizationEdge:
    from_id: str
    to_id: str
    label: str = ""
    color: str = "#666"
    directed: bool = False


@dataclass
class VisualizationStep:
    step_number: int
    description: str
    highlight_nodes: list[str] = field(default_factory=list)
    highlight_edges: list[str] = field(default_factory=list)
    code_line: int | None = None
    array_state: list[Any] | None = None
    tree_state: dict | None = None
    graph_state: dict | None = None
    stack_state: list[Any] | None = None
    queue_state: list[Any] | None = None
    linked_list_state: list[dict] | None = None


@dataclass
class VisualizationResult:
    title: str
    structure_type: str
    description: str
    steps: list[dict]
    nodes: list[dict]
    edges: list[dict]
    time_complexity: str = ""
    space_complexity: str = ""


DSA_TEMPLATES = {
    "bubble_sort": {
        "title": "Bubble Sort",
        "code": """def bubble_sort(arr):
    n = len(arr)
    for i in range(n):
        for j in range(0, n-i-1):
            if arr[j] > arr[j+1]:
                arr[j], arr[j+1] = arr[j+1], arr[j]
    return arr""",
        "description": "Bubble Sort repeatedly steps through the list, compares adjacent elements, and swaps them if they are in the wrong order.",
        "time": "O(n\u00b2)",
        "space": "O(1)",
    },
    "selection_sort": {
        "title": "Selection Sort",
        "code": """def selection_sort(arr):
    n = len(arr)
    for i in range(n):
        min_idx = i
        for j in range(i+1, n):
            if arr[j] < arr[min_idx]:
                min_idx = j
        arr[i], arr[min_idx] = arr[min_idx], arr[i]
    return arr""",
        "description": "Selection Sort divides the list into sorted and unsorted regions, repeatedly selecting the smallest element from the unsorted region.",
        "time": "O(n\u00b2)",
        "space": "O(1)",
    },
    "insertion_sort": {
        "title": "Insertion Sort",
        "code": """def insertion_sort(arr):
    for i in range(1, len(arr)):
        key = arr[i]
        j = i - 1
        while j >= 0 and key < arr[j]:
            arr[j + 1] = arr[j]
            j -= 1
        arr[j + 1] = key
    return arr""",
        "description": "Insertion Sort builds the final sorted array one element at a time by repeatedly inserting elements into their correct position.",
        "time": "O(n\u00b2)",
        "space": "O(1)",
    },
    "merge_sort": {
        "title": "Merge Sort",
        "code": """def merge_sort(arr):
    if len(arr) <= 1:
        return arr
    mid = len(arr) // 2
    left = merge_sort(arr[:mid])
    right = merge_sort(arr[mid:])
    return merge(left, right)

def merge(left, right):
    result = []
    i = j = 0
    while i < len(left) and j < len(right):
        if left[i] <= right[j]:
            result.append(left[i])
            i += 1
        else:
            result.append(right[j])
            j += 1
    result.extend(left[i:])
    result.extend(right[j:])
    return result""",
        "description": "Merge Sort is a divide-and-conquer algorithm that splits the array, recursively sorts each half, then merges the sorted halves.",
        "time": "O(n log n)",
        "space": "O(n)",
    },
    "quick_sort": {
        "title": "Quick Sort",
        "code": """def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    left = [x for x in arr if x < pivot]
    middle = [x for x in arr if x == pivot]
    right = [x for x in arr if x > pivot]
    return quick_sort(left) + middle + quick_sort(right)""",
        "description": "Quick Sort picks a pivot element and partitions the array around it, then recursively sorts the sub-arrays.",
        "time": "O(n log n)",
        "space": "O(log n)",
    },
    "binary_search": {
        "title": "Binary Search",
        "code": """def binary_search(arr, target):
    left, right = 0, len(arr) - 1
    while left <= right:
        mid = (left + right) // 2
        if arr[mid] == target:
            return mid
        elif arr[mid] < target:
            left = mid + 1
        else:
            right = mid - 1
    return -1""",
        "description": "Binary Search finds a target value in a sorted array by repeatedly dividing the search interval in half.",
        "time": "O(log n)",
        "space": "O(1)",
    },
    "linked_list": {
        "title": "Singly Linked List",
        "code": """class Node:
    def __init__(self, data):
        self.data = data
        self.next = None

class LinkedList:
    def __init__(self):
        self.head = None

    def insert_at_end(self, data):
        new_node = Node(data)
        if not self.head:
            self.head = new_node
            return
        curr = self.head
        while curr.next:
            curr = curr.next
        curr.next = new_node

    def traverse(self):
        curr = self.head
        while curr:
            print(curr.data, end=" -> ")
            curr = curr.next
        print("None")""",
        "description": "A Linked List is a linear data structure where elements are stored in nodes, each pointing to the next node.",
        "time": "O(n)",
        "space": "O(n)",
    },
    "binary_tree": {
        "title": "Binary Tree Traversals",
        "code": """class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

def inorder(root):
    if root:
        inorder(root.left)
        print(root.val, end=" ")
        inorder(root.right)

def preorder(root):
    if root:
        print(root.val, end=" ")
        preorder(root.left)
        preorder(root.right)

def postorder(root):
    if root:
        postorder(root.left)
        postorder(root.right)
        print(root.val, end=" ")""",
        "description": "A Binary Tree is a hierarchical data structure where each node has at most two children (left and right).",
        "time": "O(n)",
        "space": "O(h) where h is height",
    },
    "bfs": {
        "title": "Breadth-First Search (BFS)",
        "code": """from collections import deque

def bfs(graph, start):
    visited = set()
    queue = deque([start])
    visited.add(start)
    while queue:
        node = queue.popleft()
        print(node, end=" ")
        for neighbor in graph[node]:
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)""",
        "description": "BFS explores a graph level by level, visiting all neighbors of a node before moving to the next level.",
        "time": "O(V + E)",
        "space": "O(V)",
    },
    "dfs": {
        "title": "Depth-First Search (DFS)",
        "code": """def dfs(graph, node, visited=None):
    if visited is None:
        visited = set()
    visited.add(node)
    print(node, end=" ")
    for neighbor in graph[node]:
        if neighbor not in visited:
            dfs(graph, neighbor, visited)""",
        "description": "DFS explores a graph by going as deep as possible along each branch before backtracking.",
        "time": "O(V + E)",
        "space": "O(V)",
    },
}


def _llm_complete(prompt: str, system: str = "", temperature: float = 0.3, max_tokens: int = 512) -> str:
    payload = {
        "model": settings.slm_model,
        "prompt": f"{system}\n\n{prompt}" if system else prompt,
        "stream": False,
        "options": {"temperature": temperature, "num_predict": max_tokens},
    }
    url = f"{settings.ollama_url}/api/generate"
    try:
        data = json.dumps(payload).encode("utf-8")
        req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=60) as response:
            res = json.loads(response.read().decode("utf-8"))
            return res.get("response", "").strip()
    except Exception as e:
        print(f"LLM request failed: {e}")
        return ""


def detect_dsa_type(code: str) -> str:
    keywords = {
        "bubble_sort": ["bubble", "arr[j] > arr[j+1]", "arr[j], arr[j+1] = arr[j+1], arr[j]"],
        "selection_sort": ["selection", "min_idx", "arr[i], arr[min_idx] = arr[min_idx], arr[i]"],
        "insertion_sort": ["insertion", "key = arr[i]", "while j >= 0 and key"],
        "merge_sort": ["merge_sort", "def merge(", "left[:mid]", "right[mid:]"],
        "quick_sort": ["quick_sort", "pivot", "x < pivot", "x > pivot"],
        "binary_search": ["binary_search", "left, right = 0", "mid = (left + right)"],
        "linked_list": ["class Node", "self.data", "self.next", "class LinkedList"],
        "binary_tree": ["class TreeNode", "self.left", "self.right", "inorder", "preorder", "postorder"],
        "bfs": ["from collections import deque", "queue = deque(", "queue.popleft()"],
        "dfs": ["def dfs(", "visited.add(", "neighbor not in visited"],
        "stack": ["class Stack", "self.push", "self.pop", "self.peek", "LIFO"],
        "queue": ["class Queue", "self.enqueue", "self.dequeue", "FIFO"],
        "hash_table": ["class HashTable", "self.table", "hash(", "__getitem__"],
    }

    code_lower = code.lower()
    scores = {}
    for dsa_type, patterns in keywords.items():
        score = sum(1 for p in patterns if p.lower() in code_lower)
        if score > 0:
            scores[dsa_type] = score

    if scores:
        return max(scores, key=scores.get)
    return "general"


def get_template(template_name: str) -> dict | None:
    return DSA_TEMPLATES.get(template_name)


def list_templates() -> dict:
    return {k: {"title": v["title"], "description": v["description"], "time": v["time"], "space": v["space"]} for k, v in DSA_TEMPLATES.items()}


def generate_sorting_steps(arr: list[int], algo: str) -> list[dict]:
    steps = []
    a = list(arr)
    n = len(a)

    if algo == "bubble_sort":
        steps.append({"step_number": 0, "description": f"Initial array: {a}", "array_state": list(a), "highlight_nodes": []})
        for i in range(n):
            for j in range(0, n - i - 1):
                steps.append({"step_number": len(steps), "description": f"Comparing indices {j} and {j+1}: {a[j]} vs {a[j+1]}", "array_state": list(a), "highlight_nodes": [j, j + 1]})
                if a[j] > a[j + 1]:
                    a[j], a[j + 1] = a[j + 1], a[j]
                    steps.append({"step_number": len(steps), "description": f"Swapped {a[j+1]} and {a[j]}", "array_state": list(a), "highlight_nodes": [j, j + 1]})
        steps.append({"step_number": len(steps), "description": f"Sorted array: {a}", "array_state": list(a), "highlight_nodes": []})

    elif algo == "selection_sort":
        steps.append({"step_number": 0, "description": f"Initial array: {a}", "array_state": list(a), "highlight_nodes": []})
        for i in range(n):
            min_idx = i
            for j in range(i + 1, n):
                steps.append({"step_number": len(steps), "description": f"Comparing {a[j]} with current minimum {a[min_idx]} at index {j}", "array_state": list(a), "highlight_nodes": [j, min_idx]})
                if a[j] < a[min_idx]:
                    min_idx = j
                    steps.append({"step_number": len(steps), "description": f"New minimum found: {a[j]} at index {j}", "array_state": list(a), "highlight_nodes": [j]})
            if min_idx != i:
                a[i], a[min_idx] = a[min_idx], a[i]
                steps.append({"step_number": len(steps), "description": f"Swapped {a[i]} (index {i}) with {a[min_idx]} (index {min_idx})", "array_state": list(a), "highlight_nodes": [i, min_idx]})
        steps.append({"step_number": len(steps), "description": f"Sorted array: {a}", "array_state": list(a), "highlight_nodes": []})

    elif algo == "insertion_sort":
        steps.append({"step_number": 0, "description": f"Initial array: {a}", "array_state": list(a), "highlight_nodes": []})
        for i in range(1, n):
            key = a[i]
            j = i - 1
            steps.append({"step_number": len(steps), "description": f"Key element: {key} at index {i}", "array_state": list(a), "highlight_nodes": [i]})
            while j >= 0 and key < a[j]:
                a[j + 1] = a[j]
                steps.append({"step_number": len(steps), "description": f"Shifted {a[j]} from index {j} to {j+1}", "array_state": list(a), "highlight_nodes": [j, j + 1]})
                j -= 1
            a[j + 1] = key
            steps.append({"step_number": len(steps), "description": f"Placed key {key} at index {j+1}", "array_state": list(a), "highlight_nodes": [j + 1]})
        steps.append({"step_number": len(steps), "description": f"Sorted array: {a}", "array_state": list(a), "highlight_nodes": []})

    return steps


def generate_bst_steps(values: list[int]) -> list[dict]:
    steps = []
    nodes = []

    steps.append({"step_number": 0, "description": f"Building BST from values: {values}", "tree_state": None, "highlight_nodes": []})

    root = None
    for i, val in enumerate(values):
        if root is None:
            root = {"id": str(val), "val": val, "left": None, "right": None}
            nodes.append(root)
            steps.append({"step_number": i + 1, "description": f"Inserted root node: {val}", "tree_state": _serialize_tree(root), "highlight_nodes": [str(val)]})
        else:
            steps.append({"step_number": i + 1, "description": f"Inserting {val} into BST", "tree_state": _serialize_tree(root), "highlight_nodes": []})
            curr = root
            while True:
                if val < curr["val"]:
                    if curr["left"] is None:
                        curr["left"] = {"id": str(val), "val": val, "left": None, "right": None}
                        nodes.append(curr["left"])
                        steps.append({"step_number": len(steps), "description": f"Inserted {val} as left child of {curr['val']}", "tree_state": _serialize_tree(root), "highlight_nodes": [str(val)]})
                        break
                    curr = curr["left"]
                else:
                    if curr["right"] is None:
                        curr["right"] = {"id": str(val), "val": val, "left": None, "right": None}
                        nodes.append(curr["right"])
                        steps.append({"step_number": len(steps), "description": f"Inserted {val} as right child of {curr['val']}", "tree_state": _serialize_tree(root), "highlight_nodes": [str(val)]})
                        break
                    curr = curr["right"]

    steps.append({"step_number": len(steps), "description": f"Final BST structure", "tree_state": _serialize_tree(root), "highlight_nodes": []})
    return steps


def _serialize_tree(node) -> dict | None:
    if node is None:
        return None
    return {"id": node["id"], "val": node["val"], "left": _serialize_tree(node["left"]), "right": _serialize_tree(node["right"])}


def generate_linked_list_steps(values: list[Any]) -> list[dict]:
    steps = []
    nodes_data = [{"id": f"n{i}", "value": str(v), "next": f"n{i+1}" if i < len(values) - 1 else None} for i, v in enumerate(values)]
    steps.append({
        "step_number": 0,
        "description": f"Linked list with values: {values}",
        "linked_list_state": nodes_data,
        "highlight_nodes": [],
    })
    for i in range(len(nodes_data)):
        steps.append({
            "step_number": i + 1,
            "description": f"Traversing: node {i} = {values[i]}",
            "linked_list_state": nodes_data,
            "highlight_nodes": [f"n{i}"],
        })
    return steps


def _analyze_with_llm(code: str, language: str) -> dict:
    prompt = (
        "You are a DSA expert. Analyze the following code and provide:\n"
        "1. STRUCTURE_TYPE: What DSA structure/algorithm is this? (e.g., 'Linked List', 'Binary Tree', 'Bubble Sort', 'BFS')\n"
        "2. EXPLANATION: Step-by-step explanation of how this code works\n"
        "3. TIME_COMPLEXITY: Big O time complexity\n"
        "4. SPACE_COMPLEXITY: Big O space complexity\n"
        "5. SUMMARY: One-line summary\n\n"
        f"Language: {language}\n\n"
        f"Code:\n```{language}\n{code}\n```\n\n"
        "Format your response exactly as:\n"
        "STRUCTURE_TYPE: <type>\n"
        "EXPLANATION:\n<step-by-step explanation>\n"
        "TIME_COMPLEXITY: <complexity>\n"
        "SPACE_COMPLEXITY: <complexity>\n"
        "SUMMARY: <summary>"
    )
    response = _llm_complete(prompt, system="You are a precise DSA code analyst.", max_tokens=1024)
    result = {
        "structure_type": "Unknown",
        "explanation": "",
        "time_complexity": "",
        "space_complexity": "",
        "summary": "",
    }
    if not response:
        return result

    if "STRUCTURE_TYPE:" in response:
        parts = response.split("STRUCTURE_TYPE:", 1)
        rest = parts[1].strip()
        for key in ["EXPLANATION:", "TIME_COMPLEXITY:", "SPACE_COMPLEXITY:", "SUMMARY:"]:
            if key in rest:
                val, rest = rest.split(key, 1)
                k = key.lower().rstrip(":").replace("_", "_")
                if key == "EXPLANATION:":
                    result["explanation"] = val.strip()
                elif key == "TIME_COMPLEXITY:":
                    result["time_complexity"] = val.strip()
                elif key == "SPACE_COMPLEXITY:":
                    result["space_complexity"] = val.strip()
                elif key == "SUMMARY:":
                    result["summary"] = val.strip()
        if not result["structure_type"] or result["structure_type"] == "Unknown":
            result["structure_type"] = parts[1].split("\n")[0].strip()
    else:
        result["summary"] = response[:200]

    return result


def analyze_code(code: str, language: str = "auto") -> dict:
    dsa_type = detect_dsa_type(code)
    template = DSA_TEMPLATES.get(dsa_type)

    result = {
        "code": code,
        "language": language,
        "dsa_type": dsa_type,
        "title": template["title"] if template else dsa_type.replace("_", " ").title(),
        "description": "",
        "time_complexity": template["time"] if template else "",
        "space_complexity": template["space"] if template else "",
        "steps": [],
        "nodes": [],
        "edges": [],
        "summary": "",
        "explanation": "",
    }

    llm_result = _analyze_with_llm(code, language)
    if template:
        result["description"] = template["description"]
    result["explanation"] = llm_result.get("explanation", "") or result["description"]
    result["summary"] = llm_result.get("summary", "") or f"This code implements a {result['title']}."
    if llm_result.get("time_complexity"):
        result["time_complexity"] = llm_result["time_complexity"]
    if llm_result.get("space_complexity"):
        result["space_complexity"] = llm_result["space_complexity"]

    sample_data = [64, 34, 25, 12, 22, 11, 90]
    if dsa_type in ("bubble_sort", "selection_sort", "insertion_sort"):
        result["steps"] = generate_sorting_steps(sample_data, dsa_type)
    elif dsa_type == "binary_tree":
        result["steps"] = generate_bst_steps([10, 5, 15, 3, 7, 12, 18])
    elif dsa_type == "linked_list":
        result["steps"] = generate_linked_list_steps(["A", "B", "C", "D"])
    elif dsa_type == "quick_sort":
        result["steps"] = generate_sorting_steps([64, 34, 25, 12, 22, 11, 90], "insertion_sort")
    elif dsa_type == "merge_sort":
        result["steps"] = generate_sorting_steps([64, 34, 25, 12, 22, 11, 90], "bubble_sort")
    elif dsa_type == "binary_search":
        steps = generate_sorting_steps(sorted(sample_data), "insertion_sort")
        for s in steps:
            s["highlight_nodes"] = []
        steps.insert(0, {"step_number": 0, "description": f"Sorted array for binary search: {sorted(sample_data)}", "array_state": sorted(sample_data), "highlight_nodes": []})
        target = 22
        arr = sorted(sample_data)
        left, right = 0, len(arr) - 1
        step_num = 1
        while left <= right:
            mid = (left + right) // 2
            steps.append({"step_number": step_num, "description": f"Searching for {target}: left={left} (val={arr[left]}), right={right} (val={arr[right]}), mid={mid} (val={arr[mid]})", "array_state": arr, "highlight_nodes": [left, right, mid]})
            step_num += 1
            if arr[mid] == target:
                steps.append({"step_number": step_num, "description": f"Found {target} at index {mid}!", "array_state": arr, "highlight_nodes": [mid]})
                break
            elif arr[mid] < target:
                left = mid + 1
            else:
                right = mid - 1
        result["steps"] = steps

    return result


def visualize_template(template_name: str) -> dict | None:
    template = DSA_TEMPLATES.get(template_name)
    if not template:
        return None
    return analyze_code(template["code"], "python")


def explain_code_step_by_step(code: str, language: str = "auto") -> dict:
    dsa_type = detect_dsa_type(code)
    template = DSA_TEMPLATES.get(dsa_type)

    if not _llm_complete("ping", "say pong"):
        explanation = _basic_explain(code, dsa_type)
        summary = f"This code implements a {dsa_type.replace('_', ' ').title() if dsa_type != 'general' else 'general algorithm'}."
    else:
        llm_result = _analyze_with_llm(code, language)
        explanation = llm_result.get("explanation", "")
        summary = llm_result.get("summary", "")

    if not summary:
        summary = f"Code implements {template['title'] if template else dsa_type.replace('_', ' ').title()}."

    return {
        "language": language,
        "dsa_type": dsa_type,
        "title": template["title"] if template else dsa_type.replace("_", " ").title(),
        "explanation": explanation,
        "summary": summary,
        "time_complexity": template["time"] if template else "",
        "space_complexity": template["space"] if template else "",
    }


def _basic_explain(code: str, dsa_type: str) -> str:
    explanations = {
        "bubble_sort": "1. Start at the beginning of the array.\n2. Compare adjacent elements.\n3. If they are in the wrong order, swap them.\n4. Move to the next pair and repeat.\n5. After each pass, the largest element 'bubbles up' to its correct position.\n6. Repeat until no swaps are needed.",
        "selection_sort": "1. Find the smallest element in the unsorted portion.\n2. Swap it with the first unsorted element.\n3. Move the boundary between sorted and unsorted regions.\n4. Repeat until all elements are sorted.",
        "insertion_sort": "1. Start with the second element (key).\n2. Compare key with elements to its left.\n3. Shift larger elements right to make space.\n4. Insert the key in its correct position.\n5. Move to the next element and repeat.",
        "linked_list": "1. Each node stores data and a pointer to the next node.\n2. The head points to the first node.\n3. Traversal starts at head and follows 'next' pointers.\n4. Insertion creates a new node and updates pointers.\n5. Deletion bypasses the target node by updating the previous node's pointer.",
        "binary_tree": "1. Each node has at most two children (left and right).\n2. Inorder: left -> root -> right\n3. Preorder: root -> left -> right\n4. Postorder: left -> right -> root\n5. BST property: left < root < right.",
        "bfs": "1. Start from the source node.\n2. Visit all immediate neighbors first.\n3. Use a queue to track nodes to explore.\n4. Mark visited nodes to avoid cycles.\n5. Continue until queue is empty.",
        "dfs": "1. Start from the source node.\n2. Mark it as visited.\n3. Recursively visit each unvisited neighbor.\n4. Backtrack when no unvisited neighbors remain.\n5. Uses a stack (implicitly via recursion).",
    }
    return explanations.get(dsa_type, "Analyze the code structure to understand its logic.")