"""Controller for managing research loop state transitions.

The Controller manages:
- Budget tracking (tool calls, tokens, time)
- Model tier escalation (fast → standard → powerful)
- Question queue management
- Stopping conditions
- State transitions based on Critic decisions

Team values enforced:
- V4: Cost awareness - cheap models first, escalate only when justified
"""

import time
from typing import Optional

from rich.console import Console

from emdash.agent.research.state import (
    ResearchState,
    ResearchQuestion,
    Critique,
    IterationResult,
    FollowUpQuestion,
)


# Model tier definitions (V4: Cost awareness)
MODEL_TIERS = {
    "fast": "claude-haiku-4-5",
    "standard": "claude-sonnet-4",
    "powerful": "claude-sonnet-4",  # Could use opus when available
}


class ResearchController:
    """Manages research loop: budgets, escalation, stopping.
    
    The Controller is the "brain" that coordinates the research loop.
    It decides what to do based on Critic decisions and resource constraints.
    """
    
    def __init__(
        self,
        verbose: bool = True,
    ):
        """Initialize the controller.
        
        Args:
            verbose: Whether to print progress
        """
        self.verbose = verbose
        self.console = Console()
        self.start_time: Optional[float] = None
    
    def start_research(self, state: ResearchState) -> None:
        """Mark research start time.
        
        Args:
            state: Research state
        """
        self.start_time = time.time()
        
        # Initialize remaining budget if not set
        if not state.remaining_budget:
            state.remaining_budget = dict(state.plan.budgets)
        
        # Initialize question queue if not set
        if not state.question_queue:
            state.question_queue = [
                q.qid for q in state.plan.get_p0_questions()
            ] + [
                q.qid for q in state.plan.get_p1_questions()
            ] + [
                q.qid for q in state.plan.get_p2_questions()
            ]
    
    def select_questions(
        self,
        state: ResearchState,
        k: int = 3,
    ) -> list[ResearchQuestion]:
        """Select next questions to investigate.
        
        Prioritizes:
        1. P0 questions that haven't been answered
        2. Gap-related follow-ups from Critic
        3. P1 questions
        4. P2 questions (if budget allows)
        
        Args:
            state: Current research state
            k: Maximum number of questions to select
            
        Returns:
            List of questions to investigate
        """
        selected: list[ResearchQuestion] = []
        
        # Get unanswered questions from queue
        unanswered = [
            qid for qid in state.question_queue
            if qid not in state.answered_questions
        ]
        
        for qid in unanswered[:k]:
            question = state.get_question_by_qid(qid)
            if question:
                selected.append(question)
        
        # If we have follow-up questions from last critique, add them
        if state.history:
            last_critique = state.history[-1].critique
            for follow_up in last_critique.follow_up_questions[:k - len(selected)]:
                # Create a synthetic question from follow-up
                if follow_up.qid and follow_up.qid not in state.answered_questions:
                    question = state.get_question_by_qid(follow_up.qid)
                    if question and question not in selected:
                        selected.append(question)
                elif len(selected) < k:
                    # Create new question for follow-up
                    synthetic = ResearchQuestion(
                        qid=f"F{len(state.history)}.{len(selected)}",
                        question=follow_up.question,
                        priority="P1",
                        success_criteria=[follow_up.why],
                        suggested_tools=follow_up.suggested_tools,
                        deliverable="Implementation",
                    )
                    selected.append(synthetic)
        
        if self.verbose and selected:
            self.console.print(f"[dim]Selected {len(selected)} questions to investigate[/dim]")
        
        return selected
    
    def update_context(
        self,
        state: ResearchState,
        evidence: list,
        claims: list,
    ) -> dict:
        """Update shared context with new findings.
        
        Args:
            state: Research state
            evidence: New evidence items
            claims: New claims
            
        Returns:
            Updated context dict
        """
        context = dict(state.context)
        
        # Track discovered entities
        entities = set(context.get("discovered_entities", []))
        for e in evidence:
            entities.update(e.entities)
        context["discovered_entities"] = list(entities)[:100]  # Limit
        
        # Track key findings
        findings = context.get("key_findings", [])
        for claim in claims:
            if claim.confidence >= 2:
                findings.append({
                    "claim_id": claim.id,
                    "statement": claim.statement[:200],
                    "evidence": claim.evidence_ids,
                })
        context["key_findings"] = findings[-20:]  # Keep recent
        
        # Track answered questions
        if claims:
            # Mark questions as answered if we have high-confidence claims
            for question in state.plan.questions:
                if question.qid not in state.answered_questions:
                    # Check if any claim addresses this question
                    q_keywords = set(question.question.lower().split())
                    for claim in claims:
                        if claim.confidence >= 2:
                            c_keywords = set(claim.statement.lower().split())
                            if len(q_keywords & c_keywords) >= 3:
                                state.answered_questions.add(question.qid)
                                break
        
        return context
    
    def should_stop(self, state: ResearchState) -> tuple[bool, str]:
        """Check if research should stop.
        
        Stopping conditions:
        1. Max iterations reached
        2. Budget exhausted (>90% tool calls used)
        3. Time limit exceeded
        4. All questions answered
        
        Args:
            state: Research state
            
        Returns:
            Tuple of (should_stop, reason)
        """
        # Check max iterations
        if state.iteration >= state.plan.max_iterations:
            return True, "Max iterations reached"
        
        # Check tool call budget
        tool_budget = state.plan.budgets.get("tool_calls", 50)
        tool_remaining = state.remaining_budget.get("tool_calls", tool_budget)
        if tool_remaining <= 0:
            return True, "Tool call budget exhausted"
        if tool_remaining < tool_budget * 0.1:
            # Less than 10% remaining
            return True, "Tool call budget nearly exhausted"
        
        # Check time limit
        if self.start_time:
            elapsed = time.time() - self.start_time
            time_limit = state.plan.budgets.get("time_s", 300)
            if elapsed >= time_limit:
                return True, f"Time limit ({time_limit}s) exceeded"
        
        # Check if all P0 questions answered
        p0_qids = {q.qid for q in state.plan.get_p0_questions()}
        if p0_qids and p0_qids.issubset(state.answered_questions):
            # All P0 answered - check P1
            p1_qids = {q.qid for q in state.plan.get_p1_questions()}
            if p1_qids.issubset(state.answered_questions):
                return True, "All P0 and P1 questions answered"
        
        return False, ""
    
    def escalate_models_or_budget(
        self,
        state: ResearchState,
        critique: Critique,
    ) -> None:
        """Upgrade model tier if Critic requested ESCALATE.
        
        V4: Cost awareness - only escalate when justified.
        
        Args:
            state: Research state
            critique: Critique with ESCALATE decision
        """
        current_tier = state.current_model_tier
        
        if current_tier == "fast":
            state.current_model_tier = "standard"
            if self.verbose:
                self.console.print(
                    f"[magenta]Escalating model: fast → standard[/magenta]"
                )
        elif current_tier == "standard":
            state.current_model_tier = "powerful"
            if self.verbose:
                self.console.print(
                    f"[magenta]Escalating model: standard → powerful[/magenta]"
                )
        else:
            # Already at max tier, increase budget instead
            current_budget = state.remaining_budget.get("tool_calls", 0)
            state.remaining_budget["tool_calls"] = current_budget + 20
            if self.verbose:
                self.console.print(
                    f"[magenta]Increasing tool budget by 20 (already at powerful tier)[/magenta]"
                )
    
    def rollback_or_restart(
        self,
        state: ResearchState,
        critique: Critique,
    ) -> None:
        """Handle REJECT: fix issues or restart iteration.
        
        Args:
            state: Research state
            critique: Critique with REJECT decision
        """
        if self.verbose:
            self.console.print("[red]Research rejected. Attempting recovery...[/red]")
            for issue in critique.must_fix:
                self.console.print(f"  [red]- {issue}[/red]")
        
        # Reset the last iteration's claims (they're invalid)
        if state.history:
            last_result = state.history[-1]
            # Keep evidence but mark claims as needing rework
            state.context["rejected_claims"] = [c.id for c in last_result.claims]
        
        # Add must_fix items as new questions
        for issue in critique.must_fix[:3]:
            state.question_queue.append(f"FIX_{len(state.question_queue)}")
    
    def enqueue_followups(
        self,
        state: ResearchState,
        follow_ups: list[FollowUpQuestion],
    ) -> None:
        """Add follow-up questions to queue.
        
        Args:
            state: Research state
            follow_ups: Follow-up questions from Critic
        """
        for i, follow_up in enumerate(follow_ups[:5]):
            qid = follow_up.qid or f"FOLLOWUP_{state.iteration}_{i}"
            if qid not in state.question_queue:
                state.question_queue.append(qid)
                
                # Add to plan if it's a new question
                if not state.get_question_by_qid(qid):
                    new_question = ResearchQuestion(
                        qid=qid,
                        question=follow_up.question,
                        priority="P1",
                        success_criteria=[follow_up.why],
                        suggested_tools=follow_up.suggested_tools,
                        deliverable="Implementation",
                    )
                    state.plan.questions.append(new_question)
        
        if self.verbose and follow_ups:
            self.console.print(f"[dim]Enqueued {len(follow_ups)} follow-up questions[/dim]")
    
    def record_iteration(
        self,
        state: ResearchState,
        evidence: list,
        claims: list,
        gaps: list,
        critique: Critique,
    ) -> None:
        """Record an iteration result.
        
        Args:
            state: Research state
            evidence: Evidence collected
            claims: Claims proposed
            gaps: Gaps identified
            critique: Critic's evaluation
        """
        result = IterationResult(
            iteration=state.iteration,
            evidence=evidence,
            claims=claims,
            gaps=gaps,
            critique=critique,
            model_tier=state.current_model_tier,
        )
        
        state.history.append(result)
        
        # Update budget
        tool_calls_used = len(evidence)
        state.remaining_budget["tool_calls"] = (
            state.remaining_budget.get("tool_calls", 50) - tool_calls_used
        )
        
        # Increment iteration
        state.iteration += 1
    
    def get_model_for_tier(self, tier: str) -> str:
        """Get model name for a tier.
        
        Args:
            tier: Model tier (fast/standard/powerful)
            
        Returns:
            Model name string
        """
        return MODEL_TIERS.get(tier, MODEL_TIERS["fast"])
    
    def handle_budget_warning(self, state: ResearchState) -> bool:
        """Handle budget warning when >80% used.
        
        When budget is running low:
        - Drop P2 questions
        - Narrow P1 scope
        
        Args:
            state: Research state
            
        Returns:
            True if adjustments made
        """
        budget_used = state.budget_used_percent()
        
        if budget_used > 80:
            if self.verbose:
                self.console.print(
                    f"[yellow]Warning: {budget_used:.0f}% of tool budget used[/yellow]"
                )
            
            # Drop P2 questions
            p2_qids = {q.qid for q in state.plan.get_p2_questions()}
            state.question_queue = [
                qid for qid in state.question_queue
                if qid not in p2_qids or qid in state.answered_questions
            ]
            
            if self.verbose:
                self.console.print("[yellow]Dropped P2 questions to conserve budget[/yellow]")
            
            return True
        
        return False

