Atomizer/optimization_engine/interview/study_interview.py

"""
Study Interview Engine

Main orchestrator for the interview process.
Coordinates question flow, state management, validation, and blueprint generation.
"""

from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional, Literal
import uuid

from .interview_state import (
    InterviewState,
    InterviewPhase,
    InterviewStateManager,
    AnsweredQuestion,
    LogEntry,
)
from .question_engine import QuestionEngine, Question
from .interview_presenter import InterviewPresenter, ClaudePresenter
from .engineering_validator import EngineeringValidator, ValidationResult, AntiPattern
from .interview_intelligence import InterviewIntelligence
from .study_blueprint import StudyBlueprint, BlueprintBuilder


@dataclass
class InterviewSession:
    """Represents an active interview session."""
    session_id: str
    study_name: str
    study_path: Path
    started_at: datetime
    current_phase: InterviewPhase
    introspection: Dict[str, Any]
    is_complete: bool = False
    is_resumed: bool = False


@dataclass
class NextAction:
    """What should happen after processing an answer."""
    action_type: Literal["ask_question", "show_summary", "validate", "generate", "error", "confirm_warning"]
    question: Optional[Question] = None
    message: Optional[str] = None
    warnings: List[str] = field(default_factory=list)
    blueprint: Optional[StudyBlueprint] = None
    anti_patterns: List[AntiPattern] = field(default_factory=list)


class StudyInterviewEngine:
    """
    Main orchestrator for study interviews.

    Manages the complete interview lifecycle:
    1. Start or resume interview
    2. Present questions via presenter
    3. Process answers with validation
    4. Generate blueprint for review
    5. Handle modifications
    6. Coordinate study generation
    """

    def __init__(
        self,
        study_path: Path,
        presenter: Optional[InterviewPresenter] = None
    ):
        """
        Initialize interview engine.

        Args:
            study_path: Path to the study directory
            presenter: Presentation layer (defaults to ClaudePresenter)
        """
        self.study_path = Path(study_path)
        self.presenter = presenter or ClaudePresenter()
        self.state_manager = InterviewStateManager(self.study_path)
        self.question_engine = QuestionEngine()
        self.validator = EngineeringValidator()
        self.intelligence = InterviewIntelligence()
        self.blueprint_builder = BlueprintBuilder()

        # Current state
        self.state: Optional[InterviewState] = None
        self.introspection: Dict[str, Any] = {}
        self.current_question: Optional[Question] = None
        self.session: Optional[InterviewSession] = None

        # Estimated questions (for progress)
        self.estimated_total_questions = 12  # Will be updated based on complexity

    def start_interview(
        self,
        study_name: str,
        model_path: Optional[Path] = None,
        introspection: Optional[Dict[str, Any]] = None
    ) -> InterviewSession:
        """
        Start a new interview or resume existing one.

        Args:
            study_name: Name for the study
            model_path: Path to the NX model (optional)
            introspection: Pre-computed introspection results (optional)

        Returns:
            InterviewSession representing the active interview
        """
        # Check for existing state
        existing_state = self.state_manager.load_state()

        if existing_state and not existing_state.is_complete():
            # Resume existing interview
            self.state = existing_state
            self.introspection = existing_state.introspection

            self.session = InterviewSession(
                session_id=existing_state.session_id,
                study_name=existing_state.study_name,
                study_path=self.study_path,
                started_at=datetime.fromisoformat(existing_state.started_at),
                current_phase=existing_state.get_phase(),
                introspection=self.introspection,
                is_resumed=True,
            )

            return self.session

        # Start new interview
        self.state = InterviewState(
            session_id=str(uuid.uuid4()),
            study_name=study_name,
            study_path=str(self.study_path),
            current_phase=InterviewPhase.INTROSPECTION.value,
        )

        # Store introspection if provided
        if introspection:
            self.introspection = introspection
            self.state.introspection = introspection
            # Move to problem definition if introspection already done
            self.state.set_phase(InterviewPhase.PROBLEM_DEFINITION)

        # Save initial state
        self.state_manager.save_state(self.state)

        self.session = InterviewSession(
            session_id=self.state.session_id,
            study_name=study_name,
            study_path=self.study_path,
            started_at=datetime.now(),
            current_phase=self.state.get_phase(),
            introspection=self.introspection,
        )

        return self.session

    def get_first_question(self) -> NextAction:
        """
        Get the first question to ask.

        Returns:
            NextAction with the first question
        """
        if self.state is None:
            return NextAction(
                action_type="error",
                message="Interview not started. Call start_interview() first."
            )

        # Get next question
        next_q = self.question_engine.get_next_question(self.state, self.introspection)

        if next_q is None:
            # No questions - should not happen at start
            return NextAction(
                action_type="error",
                message="No questions available."
            )

        self.current_question = next_q

        return NextAction(
            action_type="ask_question",
            question=next_q,
            message=self.presenter.present_question(
                next_q,
                question_number=self.state.current_question_count() + 1,
                total_questions=self.estimated_total_questions,
                category_name=self._get_category_name(next_q.category),
            )
        )

    def process_answer(self, answer: str) -> NextAction:
        """
        Process user answer and determine next action.

        Args:
            answer: User's answer (natural language)

        Returns:
            NextAction indicating what to do next
        """
        if self.state is None or self.current_question is None:
            return NextAction(
                action_type="error",
                message="No active question. Call get_first_question() or get_next_question()."
            )

        question = self.current_question

        # 1. Parse answer based on question type
        parsed = self.presenter.parse_response(answer, question)

        # 2. Validate answer
        is_valid, error_msg = self.question_engine.validate_answer(parsed, question)
        if not is_valid:
            return NextAction(
                action_type="error",
                message=f"Invalid answer: {error_msg}",
                question=question,  # Re-ask same question
            )

        # 3. Store answer
        self._store_answer(question, answer, parsed)

        # 4. Update phase if needed
        self._update_phase(question)

        # 5. Update complexity after initial questions
        if question.category == "problem_definition":
            self._update_complexity()

        # 6. Check for warnings/anti-patterns
        anti_patterns = self.validator.detect_anti_patterns(self.state, self.introspection)
        new_warnings = [ap.description for ap in anti_patterns if ap.severity in ["error", "warning"]]

        # Filter to only new warnings
        existing_warnings = set(self.state.warnings)
        for w in new_warnings:
            if w not in existing_warnings:
                self.state.add_warning(w)

        # 7. Check if we should show anti-pattern warnings
        blocking_patterns = [ap for ap in anti_patterns if ap.severity == "error" and not ap.acknowledged]
        if blocking_patterns:
            return NextAction(
                action_type="confirm_warning",
                message=self._format_anti_pattern_warnings(blocking_patterns),
                anti_patterns=blocking_patterns,
            )

        # 8. Get next question
        next_q = self.question_engine.get_next_question(self.state, self.introspection)

        if next_q is None:
            # Interview complete - generate blueprint
            return self._finalize_interview()

        self.current_question = next_q

        return NextAction(
            action_type="ask_question",
            question=next_q,
            message=self.presenter.present_question(
                next_q,
                question_number=self.state.current_question_count() + 1,
                total_questions=self.estimated_total_questions,
                category_name=self._get_category_name(next_q.category),
            ),
            warnings=[w for w in self.state.warnings if w not in self.state.warnings_acknowledged],
        )

    def acknowledge_warnings(self, acknowledged: bool = True) -> NextAction:
        """
        Acknowledge current warnings and continue.

        Args:
            acknowledged: Whether user acknowledged warnings

        Returns:
            NextAction (continue or abort)
        """
        if not acknowledged:
            return NextAction(
                action_type="error",
                message="Interview paused. Please fix the issues and restart, or acknowledge warnings to proceed."
            )

        # Mark all current warnings as acknowledged
        for w in self.state.warnings:
            self.state.acknowledge_warning(w)

        # Continue to next question
        next_q = self.question_engine.get_next_question(self.state, self.introspection)

        if next_q is None:
            return self._finalize_interview()

        self.current_question = next_q

        return NextAction(
            action_type="ask_question",
            question=next_q,
            message=self.presenter.present_question(
                next_q,
                question_number=self.state.current_question_count() + 1,
                total_questions=self.estimated_total_questions,
                category_name=self._get_category_name(next_q.category),
            )
        )

    def generate_blueprint(self) -> StudyBlueprint:
        """
        Generate study blueprint from interview state.

        Returns:
            StudyBlueprint ready for generation
        """
        if self.state is None:
            raise ValueError("No interview state available")

        blueprint = self.blueprint_builder.from_interview_state(
            self.state,
            self.introspection
        )

        # Store in state
        self.state.blueprint = blueprint.to_dict()
        self.state_manager.save_state(self.state)

        return blueprint

    def modify_blueprint(self, changes: Dict[str, Any]) -> StudyBlueprint:
        """
        Apply what-if modifications to the blueprint.

        Args:
            changes: Dictionary of changes to apply

        Returns:
            Modified StudyBlueprint
        """
        if self.state is None or self.state.blueprint is None:
            raise ValueError("No blueprint available to modify")

        blueprint = StudyBlueprint.from_dict(self.state.blueprint)

        # Apply changes
        for key, value in changes.items():
            if key == "n_trials":
                blueprint.n_trials = int(value)
            elif key == "sampler":
                blueprint.sampler = value
            elif key == "add_constraint":
                # Handle adding constraints
                pass
            elif key == "remove_constraint":
                # Handle removing constraints
                pass
            # Add more modification types as needed

        # Re-validate
        validation_errors = blueprint.validate()
        if validation_errors:
            raise ValueError(f"Invalid modifications: {validation_errors}")

        # Update state
        self.state.blueprint = blueprint.to_dict()
        self.state_manager.save_state(self.state)

        return blueprint

    def confirm_blueprint(self) -> bool:
        """
        Confirm blueprint and mark interview as complete.

        Returns:
            True if successful
        """
        if self.state is None:
            return False

        self.state.set_phase(InterviewPhase.COMPLETE)
        self.state_manager.save_state(self.state)

        # Finalize log
        self.state_manager.finalize_log(self.state)

        return True

    def get_current_state(self) -> Optional[InterviewState]:
        """Get current interview state."""
        return self.state

    def get_progress(self) -> str:
        """Get formatted progress string."""
        if self.state is None:
            return "No active interview"

        return self.presenter.show_progress(
            self.state.current_question_count(),
            self.estimated_total_questions,
            self._get_phase_name(self.state.current_phase)
        )

    def reset_interview(self) -> None:
        """Reset interview and start fresh."""
        self.state_manager.delete_state()
        self.state = None
        self.current_question = None
        self.session = None

    # Private methods

    def _store_answer(self, question: Question, raw: str, parsed: Any) -> None:
        """Store answer in state."""
        # Create answered question record
        answered = AnsweredQuestion(
            question_id=question.id,
            answered_at=datetime.now().isoformat(),
            raw_response=raw,
            parsed_value=parsed,
        )

        self.state.add_answered_question(answered)

        # Map to answer field
        self._map_answer_to_field(question.maps_to, parsed)

        # Create log entry
        log_entry = LogEntry(
            timestamp=datetime.now(),
            question_id=question.id,
            question_text=question.text,
            answer_raw=raw,
            answer_parsed=parsed,
        )

        self.state_manager.append_log(log_entry)
        self.state_manager.save_state(self.state)

    def _map_answer_to_field(self, maps_to: str, value: Any) -> None:
        """Map parsed value to the appropriate answer field."""
        if not maps_to:
            return

        # Handle array indexing: "objectives[0].goal"
        if "[" in maps_to:
            import re
            match = re.match(r"(\w+)\[(\d+)\]\.(\w+)", maps_to)
            if match:
                array_name, idx, field = match.groups()
                idx = int(idx)

                # Ensure array exists
                if array_name not in self.state.answers:
                    self.state.answers[array_name] = []

                # Ensure element exists
                while len(self.state.answers[array_name]) <= idx:
                    self.state.answers[array_name].append({})

                self.state.answers[array_name][idx][field] = value
                return

        # Handle nested fields: "constraints.max_stress"
        if "." in maps_to:
            parts = maps_to.split(".")
            current = self.state.answers

            for part in parts[:-1]:
                if part not in current:
                    current[part] = {}
                current = current[part]

            current[parts[-1]] = value
            return

        # Simple field
        self.state.set_answer(maps_to, value)

    def _update_phase(self, question: Question) -> None:
        """Update interview phase based on question category."""
        category_to_phase = {
            "problem_definition": InterviewPhase.PROBLEM_DEFINITION,
            "objectives": InterviewPhase.OBJECTIVES,
            "constraints": InterviewPhase.CONSTRAINTS,
            "design_variables": InterviewPhase.DESIGN_VARIABLES,
            "physics_config": InterviewPhase.DESIGN_VARIABLES,
            "optimization_settings": InterviewPhase.VALIDATION,
            "validation": InterviewPhase.VALIDATION,
        }

        new_phase = category_to_phase.get(question.category)
        if new_phase and new_phase != self.state.get_phase():
            self.state.set_phase(new_phase)

    def _update_complexity(self) -> None:
        """Update complexity estimate after initial questions."""
        complexity = self.intelligence.determine_complexity(self.state, self.introspection)
        self.state.complexity = complexity

        # Adjust estimated questions
        if complexity == "simple":
            self.estimated_total_questions = 8
        elif complexity == "moderate":
            self.estimated_total_questions = 12
        else:
            self.estimated_total_questions = 16

    def _finalize_interview(self) -> NextAction:
        """Finalize interview and show summary."""
        self.state.set_phase(InterviewPhase.REVIEW)

        blueprint = self.generate_blueprint()

        return NextAction(
            action_type="show_summary",
            message=self.presenter.show_summary(blueprint),
            blueprint=blueprint,
        )

    def _format_anti_pattern_warnings(self, patterns: List[AntiPattern]) -> str:
        """Format anti-pattern warnings for display."""
        lines = ["**Issues Detected:**", ""]

        for ap in patterns:
            severity_icon = "X" if ap.severity == "error" else "!"
            lines.append(f"[{severity_icon}] **{ap.name}**")
            lines.append(f"   {ap.description}")
            if ap.fix_suggestion:
                lines.append(f"   *Suggestion*: {ap.fix_suggestion}")
            lines.append("")

        lines.append("Would you like to proceed anyway? Type **yes** to continue or **no** to go back and fix.")

        return "\n".join(lines)

    def _get_category_name(self, category: str) -> str:
        """Get human-readable category name."""
        names = {
            "problem_definition": "Problem Definition",
            "objectives": "Optimization Goals",
            "constraints": "Constraints",
            "design_variables": "Design Variables",
            "physics_config": "Physics Configuration",
            "optimization_settings": "Optimization Settings",
            "validation": "Validation",
        }
        return names.get(category, category.replace("_", " ").title())

    def _get_phase_name(self, phase: str) -> str:
        """Get human-readable phase name."""
        names = {
            "introspection": "Model Analysis",
            "problem_definition": "Problem Definition",
            "objectives": "Setting Objectives",
            "constraints": "Defining Constraints",
            "design_variables": "Selecting Variables",
            "validation": "Validation",
            "review": "Review & Confirm",
            "complete": "Complete",
        }
        return names.get(phase, phase.replace("_", " ").title())


# Convenience function for quick interview
def run_interview(
    study_path: Path,
    study_name: str,
    introspection: Optional[Dict[str, Any]] = None
) -> StudyInterviewEngine:
    """
    Create and start an interview engine.

    Args:
        study_path: Path to study directory
        study_name: Study name
        introspection: Optional introspection results

    Returns:
        Configured StudyInterviewEngine ready for use
    """
    engine = StudyInterviewEngine(study_path)
    engine.start_interview(study_name, introspection=introspection)
    return engine