import math
import random
from fsm_crawl.agent.generic_agent import GenericAgent

class ProbabilityBasedAgent(GenericAgent):
    BACK = "__BACK__"
    
    def __init__(self, memory_decay=0.7, min_dwell=0.5, max_dwell=2.0,
                 scroll_steps=5, backtrack_prob=0.1):
        self.memory_decay = memory_decay
        self.min_dwell = min_dwell
        self.max_dwell = max_dwell
        self.scroll_steps = scroll_steps
        self.backtrack_prob = backtrack_prob
        self.history = []
    
    def _backtrack_weight(self, links):
        """
        Probability weight for backtracking.
        Increases when:
        - few outgoing links exist
        - current page was already visited often
        """
        if len(self.history) < 2:
            return 0.0

        current = self.history[-1]
        visits = self.history.count(current)

        # Fewer options => more backtracking
        option_pressure = 1 / max(1, len(links))

        # Familiarity => boredom
        familiarity = 1 - math.exp(-visits * self.memory_decay)

        return self.backtrack_prob * (0.5 * option_pressure + 0.5 * familiarity)
        
    
    def choose_next(self, links, depth = 0):
        choices = []
        weights = []
        
        if links is None or len(links) == 0:
            if len(self.history) >= 2:
                return self.history[-2]
            else:
                return None
        
        # Link choices
        for link in links:
            visits = self.history.count(link)
            memory_score = math.exp(-visits * self.memory_decay)
            novelty = random.random()

            score = novelty + memory_score
            choices.append(link)
            weights.append(score)

        # Backtracking
        if len(self.history) > 1:
            back_w = self._backtrack_weight(links) * 12
            choices.append(self.BACK)
            weights.append(back_w)


        total = sum(weights)
        if total == 0:
            return random.choice(choices)

        picked = random.choices(
            choices,
            weights=[w / total for w in weights],
            k=1
        )[0]

        if picked not in [self.BACK]:
            self.history.append(picked)
        elif picked == self.BACK:
            picked = self.history[-2]
            self.history.append(picked)
        
        return picked
 
    def dwell(self):
        return random.uniform(self.min_dwell, self.max_dwell)
        
    def scroll(self) -> list[int]:
        scrolls = []

        for i in range(self.scroll_steps):
            if random.random() < 0.35 * i:
                break
            scrolls.append(random.randint(200, 800))

        return scrolls