"""MCP Server for Socratic Agent Generation System.

Exposes the Socratic workflow builder as MCP tools for Claude Desktop/Code.

Usage:
    python -m attune.socratic.mcp_server

Or add to Claude Desktop config:
    {
        "mcpServers": {
            "socratic": {
                "command": "python",
                "args": ["-m", "attune.socratic.mcp_server"],
                "env": {
                    "ANTHROPIC_API_KEY": "your-key"
                }
            }
        }
    }

Copyright 2026 Smart-AI-Memory
Licensed under the Apache License, Version 2.0
"""

import asyncio
import json
import logging
import os
import sys
from typing import Any

from .mcp_tools import (
    MCP_VERSION,
    SOCRATIC_TOOLS,
)

logger = logging.getLogger(__name__)


class SocraticMCPServer:
    """MCP Server exposing Socratic workflow builder tools."""

    def __init__(self):
        """Initialize the MCP server."""
        self._sessions: dict[str, Any] = {}
        self._builder = None
        self._storage = None
        self._llm_analyzer = None
        self._feedback_loop = None
        self._initialized = False

    async def _ensure_initialized(self):
        """Lazily initialize components."""
        if self._initialized:
            return

        try:
            from .engine import SocraticWorkflowBuilder
            from .feedback import FeedbackLoop
            from .llm_analyzer import LLMGoalAnalyzer
            from .storage import JSONFileStorage, set_default_storage

            # Initialize storage
            self._storage = JSONFileStorage()
            set_default_storage(self._storage)

            # Initialize builder
            self._builder = SocraticWorkflowBuilder()

            # Initialize LLM analyzer if API key available
            api_key = os.environ.get("ANTHROPIC_API_KEY")
            if api_key:
                self._llm_analyzer = LLMGoalAnalyzer(api_key=api_key)

            # Initialize feedback loop (uses default storage path)
            self._feedback_loop = FeedbackLoop()

            self._initialized = True
            logger.info("Socratic MCP Server initialized successfully")

        except ImportError as e:
            logger.warning(f"Some components not available: {e}")
            self._initialized = True

    async def handle_tool_call(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
        """Handle a tool call from Claude.

        Args:
            name: Tool name
            arguments: Tool arguments

        Returns:
            Tool result as dictionary

        """
        await self._ensure_initialized()

        handlers = {
            "socratic_start_session": self._handle_start_session,
            "socratic_set_goal": self._handle_set_goal,
            "socratic_get_questions": self._handle_get_questions,
            "socratic_submit_answers": self._handle_submit_answers,
            "socratic_generate_workflow": self._handle_generate_workflow,
            "socratic_list_sessions": self._handle_list_sessions,
            "socratic_get_session": self._handle_get_session,
            "socratic_list_blueprints": self._handle_list_blueprints,
            "socratic_analyze_goal": self._handle_analyze_goal,
            "socratic_recommend_agents": self._handle_recommend_agents,
        }

        handler = handlers.get(name)
        if not handler:
            return {"error": f"Unknown tool: {name}"}

        try:
            return await handler(arguments)
        except Exception as e:  # noqa: BLE001
            logger.exception(f"Error handling {name}")
            return {"error": str(e)}

    async def _handle_start_session(self, args: dict[str, Any]) -> dict[str, Any]:
        """Start a new Socratic session."""
        goal = args.get("goal")

        session = self._builder.start_session(goal)
        self._sessions[session.session_id] = session

        # Persist session
        if self._storage:
            self._storage.save_session(session)

        result = {
            "session_id": session.session_id,
            "state": session.state.value,
            "message": "Session started successfully",
        }

        if goal:
            result["goal"] = goal
            # Get domain from goal_analysis if available
            if session.goal_analysis:
                result["detected_domain"] = session.goal_analysis.domain

        return result

    async def _handle_set_goal(self, args: dict[str, Any]) -> dict[str, Any]:
        """Set goal for a session."""
        session_id = args["session_id"]
        goal = args["goal"]

        session = self._get_session(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        # Use LLM analyzer if available
        analysis = None
        if self._llm_analyzer:
            try:
                analysis = await self._llm_analyzer.analyze_goal(goal)
            except Exception as e:  # noqa: BLE001
                logger.warning(f"LLM analysis failed, using fallback: {e}")

        session = self._builder.set_goal(session, goal)
        self._sessions[session_id] = session

        if self._storage:
            self._storage.save_session(session)

        result = {
            "session_id": session_id,
            "state": session.state.value,
            "goal": goal,
        }
        # Add domain from goal_analysis if available
        if session.goal_analysis:
            result["detected_domain"] = session.goal_analysis.domain

        if analysis:
            result["llm_analysis"] = {
                "primary_domain": analysis.primary_domain,
                "confidence": analysis.confidence,
                "ambiguities": analysis.ambiguities,
                "suggested_questions": analysis.suggested_questions[:3],
            }

        return result

    async def _handle_get_questions(self, args: dict[str, Any]) -> dict[str, Any]:
        """Get clarifying questions for a session."""
        session_id = args["session_id"]

        session = self._get_session(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        form = self._builder.get_next_questions(session)
        if not form:
            return {
                "session_id": session_id,
                "state": session.state.value,
                "questions": [],
                "message": "No more questions - ready to generate workflow",
            }

        # Convert form to JSON-serializable format
        questions = []
        for field in form.fields:
            q = {
                "field_id": field.id,
                "type": field.field_type.value,
                "label": field.label,
                "required": field.validation.required if field.validation else False,
            }
            if field.options:
                # Serialize FieldOption objects
                q["options"] = [
                    {"value": opt.value, "label": opt.label, "description": opt.description}
                    for opt in field.options
                ]
            if field.placeholder:
                q["placeholder"] = field.placeholder
            if field.help_text:
                q["help_text"] = field.help_text
            if field.default is not None:
                q["default"] = field.default
            questions.append(q)

        return {
            "session_id": session_id,
            "state": session.state.value,
            "form_id": form.id,
            "form_title": form.title,
            "questions": questions,
        }

    async def _handle_submit_answers(self, args: dict[str, Any]) -> dict[str, Any]:
        """Submit answers to questions."""
        session_id = args["session_id"]
        answers = args["answers"]

        session = self._get_session(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        session = self._builder.submit_answers(session, answers)
        self._sessions[session_id] = session

        if self._storage:
            self._storage.save_session(session)

        ready = self._builder.is_ready_to_generate(session)

        return {
            "session_id": session_id,
            "state": session.state.value,
            "ready_to_generate": ready,
            "message": "Ready to generate workflow" if ready else "More questions available",
        }

    async def _handle_generate_workflow(self, args: dict[str, Any]) -> dict[str, Any]:
        """Generate workflow from session."""
        session_id = args["session_id"]

        session = self._get_session(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        if not self._builder.is_ready_to_generate(session):
            return {"error": "Session not ready for generation", "state": session.state.value}

        workflow = self._builder.generate_workflow(session)

        # Save blueprint
        if self._storage:
            self._storage.save_blueprint(workflow.blueprint)
            self._storage.save_session(session)

        # Convert to JSON-serializable format
        agents = []
        for agent in workflow.blueprint.agents:
            agents.append(
                {
                    "agent_id": agent.agent_id,
                    "name": agent.name,
                    "role": agent.role.value,
                    "description": agent.description,
                    "tools": [t.tool_id for t in agent.tools],
                },
            )

        stages = []
        for stage in workflow.blueprint.stages:
            stages.append(
                {
                    "stage_id": stage.stage_id,
                    "name": stage.name,
                    "agent_ids": stage.agent_ids,
                    "dependencies": stage.dependencies,
                },
            )

        metrics = []
        for metric in workflow.success_criteria.metrics:
            metrics.append(
                {
                    "metric_id": metric.metric_id,
                    "name": metric.name,
                    "description": metric.description,
                    "type": metric.metric_type.value,
                    "target": metric.target_value,
                },
            )

        return {
            "session_id": session_id,
            "blueprint_id": workflow.blueprint.blueprint_id,
            "workflow_name": workflow.blueprint.name,
            "agents": agents,
            "stages": stages,
            "success_metrics": metrics,
            "state": session.state.value,
        }

    async def _handle_list_sessions(self, args: dict[str, Any]) -> dict[str, Any]:
        """List all sessions."""
        status_filter = args.get("status_filter", "all")

        sessions = []

        # Get from storage
        if self._storage:
            stored = self._storage.list_sessions()
            for s in stored:
                if (
                    status_filter == "all"
                    or (status_filter == "active" and s.get("state") != "completed")
                    or (status_filter == "completed" and s.get("state") == "completed")
                ):
                    sessions.append(s)

        # Add in-memory sessions not yet persisted
        for sid, session in self._sessions.items():
            if not any(s.get("session_id") == sid for s in sessions):
                sessions.append(
                    {
                        "session_id": session.session_id,
                        "state": session.state.value,
                        "goal": session.goal,
                        "created_at": (
                            session.created_at.isoformat() if session.created_at else None
                        ),
                    },
                )

        return {"sessions": sessions, "count": len(sessions)}

    async def _handle_get_session(self, args: dict[str, Any]) -> dict[str, Any]:
        """Get session details."""
        session_id = args["session_id"]

        session = self._get_session(session_id)
        if not session:
            return {"error": f"Session not found: {session_id}"}

        result = {
            "session_id": session.session_id,
            "state": session.state.value,
            "goal": session.goal,
            "question_rounds": session.question_rounds,
            "created_at": session.created_at.isoformat() if session.created_at else None,
            "updated_at": session.updated_at.isoformat() if session.updated_at else None,
        }
        # Add domain from goal_analysis if available
        if session.goal_analysis:
            result["detected_domain"] = session.goal_analysis.domain
        return result

    async def _handle_list_blueprints(self, args: dict[str, Any]) -> dict[str, Any]:
        """List all blueprints."""
        domain_filter = args.get("domain_filter")

        blueprints = []

        if self._storage:
            stored = self._storage.list_blueprints()
            for bp in stored:
                if domain_filter:
                    domains = bp.get("domains", [])
                    if domain_filter not in domains:
                        continue
                blueprints.append(bp)

        return {"blueprints": blueprints, "count": len(blueprints)}

    async def _handle_analyze_goal(self, args: dict[str, Any]) -> dict[str, Any]:
        """Analyze a goal without starting a session."""
        goal = args["goal"]

        # Use LLM analyzer if available
        if self._llm_analyzer:
            try:
                analysis = await self._llm_analyzer.analyze_goal(goal)
                return {
                    "goal": goal,
                    "primary_domain": analysis.primary_domain,
                    "secondary_domains": analysis.secondary_domains,
                    "confidence": analysis.confidence,
                    "detected_requirements": analysis.detected_requirements,
                    "ambiguities": analysis.ambiguities,
                    "suggested_questions": analysis.suggested_questions,
                    "suggested_agents": analysis.suggested_agents,
                    "analysis_method": "llm",
                }
            except Exception as e:  # noqa: BLE001
                logger.warning(f"LLM analysis failed: {e}")

        # Fallback to keyword-based analysis
        domains = self._builder._detect_domains(goal)
        return {"goal": goal, "detected_domains": list(domains), "analysis_method": "keyword"}

    async def _handle_recommend_agents(self, args: dict[str, Any]) -> dict[str, Any]:
        """Get agent recommendations."""
        domains = args["domains"]
        languages = args.get("languages", [])
        quality_focus = args.get("quality_focus", [])

        from .feedback import AdaptiveAgentGenerator
        from .generator import AgentGenerator

        # Use adaptive generator if feedback available
        if self._feedback_loop:
            adaptive_gen = AdaptiveAgentGenerator(self._feedback_loop.collector)
            context = {"domains": domains, "languages": languages, "quality_focus": quality_focus}
            recommendations = adaptive_gen.recommend_agents(context)
        else:
            # Use basic generator
            generator = AgentGenerator()
            recommendations = []
            for domain in domains:
                templates = generator._get_templates_for_domain(domain)
                for t in templates:
                    recommendations.append({"template_id": t, "domain": domain, "confidence": 0.8})

        return {"recommendations": recommendations, "count": len(recommendations)}

    def _get_session(self, session_id: str):
        """Get session from memory or storage."""
        # Check memory first
        if session_id in self._sessions:
            return self._sessions[session_id]

        # Try loading from storage
        if self._storage:
            session = self._storage.load_session(session_id)
            if session:
                self._sessions[session_id] = session
                return session

        return None


async def run_mcp_server():
    """Run the MCP server using stdio transport."""
    server = SocraticMCPServer()

    # Read from stdin, write to stdout
    reader = asyncio.StreamReader()
    protocol = asyncio.StreamReaderProtocol(reader)
    await asyncio.get_event_loop().connect_read_pipe(lambda: protocol, sys.stdin)

    writer_transport, writer_protocol = await asyncio.get_event_loop().connect_write_pipe(
        asyncio.streams.FlowControlMixin,
        sys.stdout,
    )
    writer = asyncio.StreamWriter(
        writer_transport,
        writer_protocol,
        reader,
        asyncio.get_event_loop(),
    )

    async def send_response(response: dict):
        """Send a JSON-RPC response."""
        data = json.dumps(response) + "\n"
        writer.write(data.encode())
        await writer.drain()

    async def handle_message(message: dict):
        """Handle an incoming JSON-RPC message."""
        method = message.get("method")
        params = message.get("params", {})
        msg_id = message.get("id")

        if method == "initialize":
            # MCP initialization
            response = {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": {
                    "protocolVersion": MCP_VERSION,
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": "socratic-workflow-builder", "version": "1.0.0"},
                },
            }
            await send_response(response)

        elif method == "tools/list":
            # List available tools
            response = {"jsonrpc": "2.0", "id": msg_id, "result": {"tools": SOCRATIC_TOOLS}}
            await send_response(response)

        elif method == "tools/call":
            # Execute a tool
            tool_name = params.get("name")
            tool_args = params.get("arguments", {})

            result = await server.handle_tool_call(tool_name, tool_args)

            response = {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": {"content": [{"type": "text", "text": json.dumps(result, indent=2)}]},
            }
            await send_response(response)

        elif method == "notifications/initialized":
            # Client initialized notification - no response needed
            pass

        else:
            # Unknown method
            response = {
                "jsonrpc": "2.0",
                "id": msg_id,
                "error": {"code": -32601, "message": f"Method not found: {method}"},
            }
            await send_response(response)

    # Main message loop
    logger.info("Socratic MCP Server starting...")

    while True:
        try:
            line = await reader.readline()
            if not line:
                break

            message = json.loads(line.decode().strip())
            await handle_message(message)

        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON: {e}")
        except Exception as e:  # noqa: BLE001
            logger.exception(f"Error processing message: {e}")


def main():
    """Entry point for MCP server."""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        stream=sys.stderr,  # Log to stderr, MCP uses stdout
    )

    try:
        asyncio.run(run_mcp_server())
    except KeyboardInterrupt:
        logger.info("Server shutting down...")


if __name__ == "__main__":
    main()
