"""
Multi-agent architecture for advanced CRM query planning and execution.
Provides deterministic planning with OpenAI API and structured execution.
"""

import os
import json
from typing import Dict, Any, List, Optional
from pydantic import BaseModel, Field
from enum import Enum


class PlanStep(BaseModel):
    """A single step in an execution plan."""
    step_number: int
    tool: str = Field(description="Tool to use: search_customers, search_orders, search_tickets, submit_answer")
    arguments: Dict[str, Any] = Field(description="Arguments for the tool")
    rationale: str = Field(description="Why this step is needed")
    expected_output: str = Field(description="What result we expect")


class Plan(BaseModel):
    """Complete execution plan generated by planner agent."""
    task_id: str
    description: str
    total_steps: int
    steps: List[PlanStep]
    reasoning: str = Field(description="Overall reasoning for the plan")


class AgentRole(str, Enum):
    """Agent roles in the multi-agent system."""
    PLANNER = "planner"
    EXECUTOR = "executor"
    COORDINATOR = "coordinator"


class PlannerAgent:
    """
    Planner agent that generates structured execution plans.
    Uses OpenAI API with temperature=0 for deterministic output.
    """

    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize planner agent.
        
        Args:
            api_key: OpenAI API key (defaults to OPENAI_API_KEY env var)
        """
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        self.model = "gpt-3.5-turbo"
        self.temperature = 0.0  # Deterministic

    def generate_plan(
        self,
        task_id: str,
        task_description: str,
        tables_summary: Dict[str, Any],
        max_steps: int = 15,
    ) -> Plan:
        """
        Generate execution plan for a task.
        
        Args:
            task_id: Task ID
            task_description: Task description
            tables_summary: Summary of available data
            max_steps: Maximum steps allowed
        
        Returns:
            Plan object with structured steps
        """
        try:
            import openai
        except ImportError:
            raise ImportError("openai not installed: pip install openai")

        if not self.api_key:
            raise ValueError("OPENAI_API_KEY not set")

        openai.api_key = self.api_key

        prompt = f"""You are an expert database analyst planning a CRM query task.
        
Task ID: {task_id}
Task Description: {task_description}
Database Summary: {json.dumps(tables_summary, indent=2)}
Max Steps: {max_steps}

Available tools:
1. search_customers: Filter by customer_id, name, email, tier, phone
2. search_orders: Filter by order_id, customer_id, product, status
3. search_tickets: Filter by ticket_id, customer_id, priority, status
4. submit_answer: Submit customer_ids list

Generate a detailed JSON plan with the following structure:
{{
  "task_id": "{task_id}",
  "description": "Brief description",
  "total_steps": <number>,
  "reasoning": "Overall strategy",
  "steps": [
    {{
      "step_number": 1,
      "tool": "search_customers|search_orders|search_tickets|submit_answer",
      "arguments": {{"key": "value"}},
      "rationale": "Why this step",
      "expected_output": "What we expect to find"
    }},
    ...
  ]
}}

Important: Output ONLY valid JSON, no other text. Keep total_steps <= {max_steps}."""

        try:
            response = openai.ChatCompletion.create(
                model=self.model,
                messages=[
                    {
                        "role": "system",
                        "content": "You are a JSON-only response generator. Always respond with valid JSON only."
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                temperature=self.temperature,
                max_tokens=1500
            )

            response_text = response.choices[0].message.content.strip()

            # Parse JSON response
            try:
                plan_data = json.loads(response_text)
            except json.JSONDecodeError:
                # Try to extract JSON from response
                import re
                json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                if json_match:
                    plan_data = json.loads(json_match.group())
                else:
                    raise ValueError("Could not parse plan JSON")

            # Validate and construct Plan
            plan = Plan(
                task_id=plan_data.get("task_id", task_id),
                description=plan_data.get("description", task_description),
                total_steps=plan_data.get("total_steps", 1),
                reasoning=plan_data.get("reasoning", ""),
                steps=[
                    PlanStep(
                        step_number=step["step_number"],
                        tool=step["tool"],
                        arguments=step["arguments"],
                        rationale=step["rationale"],
                        expected_output=step["expected_output"]
                    )
                    for step in plan_data.get("steps", [])
                ]
            )

            return plan

        except Exception as e:
            # Fallback: generate simple plan
            return self._generate_fallback_plan(task_id, task_description)

    def _generate_fallback_plan(self, task_id: str, task_description: str) -> Plan:
        """Generate a simple fallback plan when API fails."""
        return Plan(
            task_id=task_id,
            description=task_description,
            total_steps=2,
            reasoning="Fallback plan: first search, then submit",
            steps=[
                PlanStep(
                    step_number=1,
                    tool="search_customers",
                    arguments={"tier": "Gold"},
                    rationale="Initial search to explore data",
                    expected_output="List of customers"
                ),
                PlanStep(
                    step_number=2,
                    tool="submit_answer",
                    arguments={"customer_ids": []},
                    rationale="Submit findings",
                    expected_output="Task completion"
                )
            ]
        )


class ExecutorAgent:
    """
    Executor agent that follows a plan to interact with the environment.
    Tracks memory usage and execution quality.
    """

    def __init__(self):
        """Initialize executor agent."""
        self.execution_history: List[Dict[str, Any]] = []
        self.retrieved_entities: Dict[str, List[Dict]] = {
            "customers": [],
            "orders": [],
            "tickets": []
        }
        self.memory_hits: int = 0
        self.memory_misses: int = 0

    def reset(self) -> None:
        """Reset executor state."""
        self.execution_history = []
        self.retrieved_entities = {
            "customers": [],
            "orders": [],
            "tickets": []
        }
        self.memory_hits = 0
        self.memory_misses = 0

    def execute_plan(
        self,
        plan: Plan,
        env: Any,  # CRMQueryEnv instance
    ) -> Dict[str, Any]:
        """
        Execute plan steps against environment.
        
        Args:
            plan: Plan object with steps
            env: Environment instance
        
        Returns:
            Execution result dict with scores
        """
        self.reset()
        results = {
            "plan_id": plan.task_id,
            "steps_executed": 0,
            "final_answer": None,
            "memory_efficiency": 0.0,
            "plan_adherence": 1.0,
            "execution_history": []
        }

        for step in plan.steps:
            if env.done:
                break

            # Execute step
            action = {
                "tool": step.tool,
                "arguments": step.arguments
            }

            try:
                obs, reward, done, info = env.step(action)
                results["steps_executed"] += 1

                # Track retrieved entities
                if step.tool == "search_customers":
                    data = obs.last_action_result.get("data", [])
                    self.retrieved_entities["customers"].extend(data)
                elif step.tool == "search_orders":
                    data = obs.last_action_result.get("data", [])
                    self.retrieved_entities["orders"].extend(data)
                elif step.tool == "search_tickets":
                    data = obs.last_action_result.get("data", [])
                    self.retrieved_entities["tickets"].extend(data)
                elif step.tool == "submit_answer":
                    results["final_answer"] = step.arguments.get("customer_ids", [])

                self.execution_history.append({
                    "step": step.step_number,
                    "tool": step.tool,
                    "reward": reward.value,
                    "done": done
                })

                if done:
                    break

            except Exception as e:
                results["error"] = str(e)
                break

        # Calculate memory efficiency
        total_queries = results["steps_executed"]
        if total_queries > 0:
            unique_queries = len(set(str(step.arguments) for step in plan.steps))
            results["memory_efficiency"] = unique_queries / total_queries

        return results


class Coordinator:
    """
    Coordinator manages interaction between planner and executor agents.
    Orchestrates the full multi-agent pipeline.
    """

    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize coordinator.
        
        Args:
            api_key: OpenAI API key
        """
        self.planner = PlannerAgent(api_key)
        self.executor = ExecutorAgent()
        self.execution_log: List[Dict[str, Any]] = []

    def run_pipeline(
        self,
        env: Any,  # CRMQueryEnv instance
        max_iterations: int = 1,
    ) -> Dict[str, Any]:
        """
        Run full planning and execution pipeline.
        
        Args:
            env: Environment instance
            max_iterations: Max planner iterations
        
        Returns:
            Pipeline results
        """
        if not env.current_task_id:
            raise ValueError("Environment not reset")

        from .tasks import get_task_by_id
        task = get_task_by_id(env.current_task_id)

        results = {
            "task_id": env.current_task_id,
            "iterations": 0,
            "plans": [],
            "execution_results": [],
            "final_answer": None,
            "episode_reward": 0.0,
            "success": False
        }

        for iteration in range(max_iterations):
            results["iterations"] += 1

            # Phase 1: Planner generates plan
            try:
                plan = self.planner.generate_plan(
                    task_id=env.current_task_id,
                    task_description=task.description,
                    tables_summary=env.state().tables_summary,
                    max_steps=task.max_steps - env.step_count
                )
                results["plans"].append(plan.model_dump())
            except Exception as e:
                results["plan_error"] = str(e)
                break

            # Phase 2: Executor follows plan
            try:
                exec_result = self.executor.execute_plan(plan, env)
                results["execution_results"].append(exec_result)

                if exec_result.get("final_answer"):
                    results["final_answer"] = exec_result["final_answer"]
                    results["success"] = True
                    break

            except Exception as e:
                results["execution_error"] = str(e)
                break

            # Check if done
            if env.done:
                break

        results["episode_reward"] = env.episode_reward
        results["steps_taken"] = env.step_count

        return results