"""Conversation-aware Promethian Agent with multi-turn support."""

from __future__ import annotations

import time
from datetime import datetime, timedelta
from typing import Any

from promethian.config import PromethianConfig, get_config
from promethian.core.types import (
    Skill, Tool, Plan, ExecutionResult, ExecutionStatus,
    MessageType, SatisfactionType,
    Message, SatisfactionSignal, ConversationSession,
)
from promethian.core.planner import Planner
from promethian.core.executor import Executor
from promethian.core.context_gatherer import ContextGatherer, GatheredContext
from promethian.core.message_classifier import MessageClassifier
from promethian.cache import LocalCache
from promethian.learning import ToolBuilder, SkillLearner
from promethian.utils.claude import ClaudeClient
from promethian.api.client import PromethianAPIClient


class ConversationAgent:
    """
    Conversation-aware Promethian Agent with multi-turn support.

    Unlike the basic PromethianAgent which treats each run() as isolated,
    this agent maintains conversation context across multiple messages
    and tracks user satisfaction signals.

    Usage:
        async with ConversationAgent() as agent:
            # Start a session
            await agent.start_session("Help me build a web scraper")

            # Multiple turns
            response1 = await agent.message("First, show me how to fetch a page")
            response2 = await agent.message("Now extract the links")
            response3 = await agent.message("Perfect, thanks!")

            # End session (or let context manager handle it)
            await agent.end_session()
    """

    def __init__(
        self,
        config: PromethianConfig | None = None,
        enable_context_gathering: bool = True,
    ):
        self.config = config or get_config()
        self.enable_context_gathering = enable_context_gathering

        # Initialize Claude client
        self.claude = ClaudeClient(
            api_key=self.config.anthropic_api_key,
            model=self.config.model,
            max_tokens=self.config.max_tokens,
        )

        # Initialize components
        self.context_gatherer = ContextGatherer(self.config)
        self.planner = Planner(self.config, self.claude)
        self.executor = Executor(self.config, self.claude)
        self.tool_builder = ToolBuilder(self.claude)
        self.skill_learner = SkillLearner(self.claude)
        self.message_classifier = MessageClassifier(self.claude)
        self.cache = LocalCache()

        # API client
        self._api_client: PromethianAPIClient | None = None

        # Active conversation session
        self._session: ConversationSession | None = None
        self._api_session_id: str | None = None

    async def _get_api_client(self) -> PromethianAPIClient:
        """Get or create the API client."""
        if self._api_client is None:
            if not self.config.promethian_api_key:
                raise ValueError("PROMETHIAN_API_KEY required for API access")
            self._api_client = PromethianAPIClient(
                api_url=self.config.promethian_api_url,
                api_key=self.config.promethian_api_key,
            )
        return self._api_client

    async def close(self) -> None:
        """Close the agent and clean up resources."""
        # End session if active
        if self._session:
            await self.end_session(user_initiated=False)

        await self.executor.close()
        await self.context_gatherer.close()

        if self._api_client:
            await self._api_client.close()
            self._api_client = None

    async def __aenter__(self) -> "ConversationAgent":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        await self.close()

    @property
    def has_active_session(self) -> bool:
        """Check if there's an active session."""
        return self._session is not None

    @property
    def session(self) -> ConversationSession | None:
        """Get the current session."""
        return self._session

    async def start_session(self, task_description: str) -> ConversationSession:
        """
        Start a new conversation session.

        Args:
            task_description: The initial task/goal for this session

        Returns:
            The created ConversationSession
        """
        if self._session:
            raise ValueError("Session already active. Call end_session() first.")

        # Start API session
        api = await self._get_api_client()
        api_session = await api.start_session(task_description)
        self._api_session_id = api_session.session_id

        # Create local session
        self._session = ConversationSession(
            session_id=api_session.session_id,
            task_description=task_description,
            expires_at=api_session.expires_at,
        )

        # Clear cache for new session
        self.cache.on_session_start(api_session.session_id)

        # Pre-fetch relevant skills and tools for the initial task
        await self._fetch_knowledge(task_description)

        return self._session

    async def _fetch_knowledge(self, query: str) -> None:
        """Fetch relevant skills and tools for a query."""
        import asyncio
        api = await self._get_api_client()

        # Search skills and tools in parallel
        skills_task = api.search_skills(
            query=query,
            max_results=self.config.max_retrieved_skills,
            include_content=True,
        )
        tools_task = api.search_tools(
            query=query,
            max_results=self.config.max_retrieved_tools,
            include_content=True,
        )

        (skill_results, skills_query_id), (tool_results, tools_query_id) = await asyncio.gather(
            skills_task, tools_task
        )

        for result in skill_results:
            if result.content:
                skill = result.to_skill()
                self._session.retrieved_skills[skill.id] = skill
                self.cache.put_skills([skill], query_id=skills_query_id)

        for result in tool_results:
            if result.code:
                tool = result.to_tool()
                self._session.retrieved_tools[tool.id] = tool
                self.cache.put_tools([tool], query_id=tools_query_id)

    async def message(self, user_input: str) -> str:
        """
        Process a user message within the current session.

        Args:
            user_input: The user's message

        Returns:
            The agent's response
        """
        if not self._session:
            raise ValueError("No active session. Call start_session() first.")

        message_index = len(self._session.messages)
        is_first = message_index == 0

        # Classify the message
        msg_type, confidence = await self.message_classifier.classify(
            message=user_input,
            conversation_context=self._session.get_conversation_context() if not is_first else None,
            is_first_message=is_first,
        )

        # Check if message references prior output
        last_assistant_msg = None
        for msg in reversed(self._session.messages):
            if msg.role == "assistant":
                last_assistant_msg = msg.content
                break

        references_prior = self.message_classifier.references_prior_output(
            user_input, last_assistant_msg
        )

        # Record user message
        user_message = Message(
            role="user",
            content=user_input,
            message_type=msg_type,
            references_prior=references_prior,
        )
        self._session.messages.append(user_message)

        # Record satisfaction signal if applicable
        signal = self.message_classifier.get_satisfaction_signal(msg_type, confidence)
        if signal:
            self._session.add_signal(
                signal_type=signal[0],
                message_index=message_index,
                confidence=signal[1],
            )

        # Handle special message types
        if msg_type == MessageType.CLOSE_SESSION:
            # User is done
            response = "Session complete. Let me know if you need anything else!"
            await self._record_assistant_message(response)
            await self.end_session(user_initiated=True)
            return response

        if msg_type == MessageType.POSITIVE_FEEDBACK:
            # Just acknowledge, don't execute anything
            response = "Glad I could help! Let me know if there's anything else."
            await self._record_assistant_message(response)
            return response

        # For actual requests, execute
        response = await self._execute_request(user_input, msg_type)
        await self._record_assistant_message(response)

        return response

    async def _record_assistant_message(self, content: str) -> None:
        """Record an assistant message in the session."""
        self._session.messages.append(Message(
            role="assistant",
            content=content,
        ))

    async def _execute_request(
        self,
        user_input: str,
        msg_type: MessageType,
    ) -> str:
        """Execute a user request and return the response."""
        start_time = time.time()

        # For follow-ups or refinements, we might need additional knowledge
        if msg_type in (MessageType.FOLLOW_UP, MessageType.NEW_REQUEST):
            await self._fetch_knowledge(user_input)

        # Get available skills and tools
        skills = list(self._session.retrieved_skills.values())
        tools = list(self._session.retrieved_tools.values())

        # Gather context if needed
        gathered_context: GatheredContext | None = None
        if self.enable_context_gathering:
            gathered_context = await self.context_gatherer.gather_context(
                task=user_input,
                skills=skills,
                tools=tools,
            )

        # Build full context including conversation history
        context_parts = []

        # Add conversation history
        if len(self._session.messages) > 1:
            context_parts.append("Previous conversation:\n" + self._session.get_conversation_context())

        # Add gathered context
        if gathered_context and gathered_context.summary:
            context_parts.append(gathered_context.to_prompt_context())

        full_context = "\n\n".join(context_parts) if context_parts else None

        # Create plan
        plan = await self.planner.create_plan(
            task=user_input,
            context=full_context,
            skills=skills,
            tools=tools,
        )

        # Build any required tools
        built_tools = await self._build_required_tools(plan)
        if built_tools:
            plan.relevant_tools.extend(built_tools)
            for tool in built_tools:
                self._session.retrieved_tools[tool.id] = tool

        # Execute plan
        result = await self.executor.execute_plan(plan)

        # Record execution result
        self._session.execution_history.append(result)

        # Record technical success/failure signal
        if result.status == ExecutionStatus.SUCCESS:
            self._session.add_signal(SatisfactionType.TECHNICAL_SUCCESS)
        else:
            self._session.add_signal(SatisfactionType.TECHNICAL_FAILURE)

        # Learn from result
        await self._learn_from_result(result)

        # Format response
        if result.status == ExecutionStatus.SUCCESS:
            response = result.final_output or "Task completed successfully."
        else:
            response = f"I encountered an issue: {result.error or 'Unknown error'}"

        return response

    async def _build_required_tools(self, plan: Plan) -> list[Tool]:
        """Build any tools that the plan requires."""
        built_tools: list[Tool] = []
        seen_descriptions: set[str] = set()

        for step in plan.steps:
            for tool_description in step.tools_to_build:
                if tool_description in seen_descriptions:
                    continue
                seen_descriptions.add(tool_description)

                try:
                    tool = await self.tool_builder.build_tool(
                        description=tool_description,
                        context=f"For step: {step.description}",
                    )
                    built_tools.append(tool)
                    step.tools_to_use.append(tool.id)
                except Exception:
                    pass

        return built_tools

    async def _learn_from_result(self, result: ExecutionResult) -> None:
        """
        Learn from an execution result.

        This does two things:
        1. Extracts new skills from successful novel executions
        2. Improves existing skills/tools that needed self-healing
        """
        if result.status != ExecutionStatus.SUCCESS:
            return

        # 1. Check for new skill extraction
        if await self.skill_learner.should_extract_skill(result):
            skill = await self.skill_learner.extract_skill(result)
            if skill:
                self._session.retrieved_skills[skill.id] = skill

        # 2. Improve skills/tools that needed self-healing
        await self._improve_healed_items(result)

    async def _improve_healed_items(self, result: ExecutionResult) -> None:
        """
        Improve skills/tools that needed self-healing to succeed.

        When a step uses a skill/tool but required healing attempts,
        the healing log contains what went wrong and how it was fixed.
        We use this to create improved versions.
        """
        for step_result in result.step_results:
            # Only look at steps that succeeded but needed healing
            if step_result.status != ExecutionStatus.SUCCESS:
                continue
            if step_result.attempts <= 1 or not step_result.healing_log:
                continue  # No healing needed - item worked first try

            # Build healing context from the log
            healing_context = "\n".join(step_result.healing_log)

            # Improve skills that were used in this step
            for skill_id in step_result.skills_used:
                if skill_id in self._session.retrieved_skills:
                    skill = self._session.retrieved_skills[skill_id]
                    # Don't re-improve the same skill
                    if skill_id in self._session.improved_skills:
                        continue

                    try:
                        improved = await self.skill_learner.improve_skill(
                            skill=skill,
                            new_execution=result,
                            what_changed=healing_context,
                        )
                        if improved:
                            self._session.improved_skills[improved.id] = improved
                    except Exception:
                        pass  # Don't fail if improvement fails

            # Improve tools that were used in this step
            for tool_id in step_result.tools_used:
                if tool_id in self._session.retrieved_tools:
                    tool = self._session.retrieved_tools[tool_id]
                    # Don't re-improve the same tool
                    if tool_id in self._session.improved_tools:
                        continue

                    # Extract error from healing log for tool improvement
                    # The healing log has format: "Attempt N: error\nFix: suggestion"
                    error_parts = []
                    for log_entry in step_result.healing_log:
                        if ":" in log_entry:
                            error_parts.append(log_entry.split(":", 1)[1].strip())
                    error_summary = "\n".join(error_parts) if error_parts else healing_context

                    try:
                        improved = await self.tool_builder.improve_tool(
                            tool=tool,
                            error=error_summary,
                            context=f"Step: {step_result.step_id}",
                        )
                        if improved:
                            self._session.improved_tools[improved.id] = improved
                    except Exception:
                        pass  # Don't fail if improvement fails

    async def end_session(self, user_initiated: bool = True) -> float:
        """
        End the current session and return the success score.

        Args:
            user_initiated: Whether the user explicitly ended the session
                           (vs timeout/abandon)

        Returns:
            The calculated success score (0-1)
        """
        if not self._session:
            return 0.5  # Neutral if no session

        # Add closing signal
        if user_initiated:
            self._session.add_signal(SatisfactionType.IMPLICIT_POSITIVE_CLOSE)
        else:
            self._session.add_signal(SatisfactionType.IMPLICIT_NEGATIVE_ABANDON)

        # Calculate success score
        success_score = self._session.get_success_score()

        # Build feedback for API
        items_feedback = []
        for skill_id, skill in self._session.retrieved_skills.items():
            # Check if skill was used in any execution
            was_used = any(
                skill_id in (r.step_results[0].skills_used if r.step_results else [])
                for r in self._session.execution_history
            )
            items_feedback.append({
                "item_id": skill_id,
                "item_type": "skill",
                "was_used": was_used,
                "execution_success": success_score >= 0.5 if was_used else None,
            })

        for tool_id, tool in self._session.retrieved_tools.items():
            was_used = any(
                tool_id in (r.step_results[0].tools_used if r.step_results else [])
                for r in self._session.execution_history
            )
            items_feedback.append({
                "item_id": tool_id,
                "item_type": "tool",
                "was_used": was_used,
                "execution_success": success_score >= 0.5 if was_used else None,
            })

        # Build improved skills for contribution
        new_skills = []
        for skill in self._session.improved_skills.values():
            new_skills.append({
                "name": skill.name,
                "description": skill.description,
                "content": skill.content,
                "tags": skill.tags,
                "source_task": skill.source_task,
                "execution_success": True,  # Only improved after success
                "derived_from": skill.derived_from,
            })

        # Build improved tools for contribution
        new_tools = []
        for tool in self._session.improved_tools.values():
            new_tools.append({
                "name": tool.name,
                "description": tool.description,
                "code": tool.code,
                "function_name": tool.function_name,
                "parameters": tool.parameters,
                "dependencies": tool.dependencies,
                "tags": tool.tags,
                "source_task": tool.source_task,
                "execution_success": True,
                "derived_from": tool.derived_from,
            })

        # Complete API session
        try:
            api = await self._get_api_client()
            await api.complete_session(
                session_id=self._api_session_id,
                success_score=success_score,
                items_feedback=items_feedback,
                new_skills=new_skills,
                new_tools=new_tools,
                satisfaction_signals=[s.to_dict() for s in self._session.satisfaction_signals],
                conversation_length=self._session.conversation_length,
                had_refinements=self._session.had_refinements,
                user_initiated_close=user_initiated,
            )
        except Exception:
            pass  # Don't fail if API completion fails

        # Clear cache
        self.cache.on_session_end()

        # Clear session
        session_data = self._session.to_dict()
        self._session = None
        self._api_session_id = None

        return success_score

    def get_session_summary(self) -> dict | None:
        """Get a summary of the current session."""
        if not self._session:
            return None

        return {
            "session_id": self._session.session_id,
            "task": self._session.task_description,
            "turns": self._session.conversation_length,
            "success_score": self._session.get_success_score(),
            "had_refinements": self._session.had_refinements,
            "skills_used": len(self._session.retrieved_skills),
            "tools_used": len(self._session.retrieved_tools),
            "skills_improved": len(self._session.improved_skills),
            "tools_improved": len(self._session.improved_tools),
            "executions": len(self._session.execution_history),
        }
