"""
{{ agent_name | to_pascal_case }} Agent - Google ADK Implementation

Auto-generated from SuperSpec playbook using SuperOptiX compiler.
Framework: Google ADK (Agent Development Kit)
Generated: {{ timestamp }}

SuperSpec Metadata:
  Name: {{ metadata.name }}
  Version: {{ metadata.version }}
  Description: {{ metadata.description }}
  Framework: Google ADK
"""

from typing import List, Dict, Any, Optional
import asyncio
import uuid
import json
from pathlib import Path

# Google ADK imports
try:
  from google.adk import Agent, Runner
  from google.adk.runners import InMemoryRunner
  from google.genai import types
  ADK_AVAILABLE = True
except ImportError as e:
  ADK_AVAILABLE = False
  print(f"⚠️  Google ADK not available: {e}")
  print("⚠️  Install with: pip install google-adk")
  print("⚠️  Set GOOGLE_API_KEY environment variable")

from superoptix.core.base_component import BaseComponent
from superoptix.runners.google_adk_runtime_helpers import (
  get_google_adk_rlm_config as superoptix_get_google_adk_rlm_config,
  run_agent_with_optional_rlm as superoptix_run_agent_with_optional_rlm,
)

COMPILED_SPEC_PATH = Path(__file__).resolve().parent / "{{ compiled_spec_filename }}"
def _load_compiled_spec(path: Path) -> Dict[str, Any]:
  try:
    with path.open("r", encoding="utf-8") as _spec_file:
      return json.load(_spec_file)
  except FileNotFoundError as exc:
    raise FileNotFoundError(
      f"Compiled spec file not found at {path}. Recompile this agent to regenerate pipeline artifacts."
    ) from exc
  except json.JSONDecodeError as exc:
    raise ValueError(
      f"Compiled spec file is invalid JSON at {path}. Recompile this agent to regenerate pipeline artifacts."
    ) from exc


FULL_SPEC = _load_compiled_spec(COMPILED_SPEC_PATH)


class {{ agent_name | to_pascal_case }}Component(BaseComponent):
  """
  BaseComponent wrapper for Google ADK Agent - {{ metadata.description }}
  
  This component wraps a Google ADK Agent and makes it compatible
  with SuperOptiX's Universal GEPA optimizer.
  
  Optimizable Variable: instruction (system prompt for agent behavior)
  
  Framework: Google ADK
  Model: Gemini (optimized for Google models)
  Input: Any field from input_fields
  Output: Any field from output_fields
  """
  
  def __init__(
    self,
    instruction: Optional[str] = None,
    model_config: Optional[Dict] = None,
    tools: Optional[List] = None,
    **kwargs
  ):
    """
    Initialize Google ADK agent component.
    
    Args:
      instruction: Agent instructions (optimizable by GEPA!)
      model_config: Model configuration dict
      tools: List of tools for the agent
      **kwargs: Additional configuration
    """
    # Default instruction from playbook
    default_instruction = self._build_default_instruction()
    
    super().__init__(
      name="{{ agent_name }}",
      description="""{{ metadata.description | replace('"', '\\"') }}""",
      input_fields={{ spec.input_fields | map(attribute='name') | list | tojson }},
      output_fields={{ spec.output_fields | map(attribute='name') | list | tojson }},
      variable=instruction or default_instruction,  # GEPA optimizes this!
      variable_type="instruction",
      framework="google_adk",
      config={
        "model": model_config.get("model", "gemini-2.0-flash") if model_config else "gemini-2.0-flash",
        "description": model_config.get("description", "{{ metadata.description }}") if model_config else "{{ metadata.description }}",
        "tools": tools or [],
        **kwargs
      }
    )
    
    self._agent = None
    self._runner = None
    self._rlm_config = superoptix_get_google_adk_rlm_config(FULL_SPEC)
    if self._rlm_config.get("enabled", False):
      print(
        "✅ RLM configured "
        f"(mode={self._rlm_config.get('mode', 'assist')}, "
        f"backend={self._rlm_config.get('backend', 'litellm')})"
      )
  
  def _build_default_instruction(self) -> str:
    """Build default instruction from playbook."""
    instruction_parts = []
    
{% if spec.persona.role %}
    instruction_parts.append("You are a " + {{ (spec.persona.role | replace('\n', ' ')) | tojson }})
{% endif %}
{% if spec.persona.instructions %}
    instruction_parts.append("Instructions: " + {{ (spec.persona.instructions | replace('\n', ' ')) | tojson }})
{% endif %}
{% if spec.persona.goal %}
    instruction_parts.append("Your goal is: " + {{ (spec.persona.goal | replace('\n', ' ')) | tojson }})
{% endif %}
{% if spec.persona.backstory %}
    instruction_parts.append({{ (spec.persona.backstory | replace('\n', ' ')) | tojson }})
{% endif %}
{% if spec.persona.traits %}
    instruction_parts.append("Traits: " + {{ (spec.persona.traits | join(', ') | replace('\n', ' ')) | tojson }})
{% endif %}
{% if spec.reasoning and spec.reasoning.steps %}
    instruction_parts.append("Your approach:")
{% for step in spec.reasoning.steps %}
    instruction_parts.append("  {{ loop.index }}. {{ step }}")
{% endfor %}
{% endif %}
{% if spec.persona.constraints %}
    instruction_parts.append("Constraints:")
{% for constraint in spec.persona.constraints %}
    instruction_parts.append("  - {{ constraint }}")
{% endfor %}
{% endif %}
    
    if not instruction_parts:
      return "You are a helpful AI assistant."
    
    return "\n".join(instruction_parts)
  
  def _initialize_agent(self):
    """Lazy initialization of Google ADK agent."""
    if self._agent is not None:
      return
    
    if not ADK_AVAILABLE:
      raise ImportError("Google ADK not installed. Install with: pip install google-adk")
    
    # Convert config to dict if needed
    import types as python_types
    config = vars(self.config) if isinstance(self.config, python_types.SimpleNamespace) else self.config
    if not isinstance(config, dict):
      config = {}
    
    model_str = config.get("model", "gemini-2.0-flash")
    description = config.get("description", "AI Agent")
    tools = config.get("tools", [])
    
    # Create Google ADK Agent
    self._agent = Agent(
      model=model_str,
      name="{{ agent_name }}",
      description=description,
      instruction=self.variable,  # GEPA optimizes this!
      tools=tools,
    )
    
    # Create InMemoryRunner (simplest for evaluation)
    self._runner = InMemoryRunner(
      agent=self._agent,
      app_name="superoptix_{{ agent_name }}",
    )
    
    print(f"✅ Google ADK Agent initialized with model: {model_str}")
  
  def forward(self, **inputs: Any) -> Dict[str, Any]:
    """
    Execute the Google ADK agent.
    
    Args:
      **inputs: Input fields from playbook
    
    Returns:
      Dict with output fields
    """
    # Ensure agent is initialized
    self._initialize_agent()
    
    # Get user message from inputs
{% if spec.input_fields %}
{% set first_input = spec.input_fields[0] %}
    user_message = inputs.get("{{ first_input.name }}", "")
{% else %}
    user_message = inputs.get("query", inputs.get("message", ""))
{% endif %}
    
    # Execute agent
    try:
      # Generate unique IDs
      user_id = "superoptix_user"
      session_id = str(uuid.uuid4())
      
      # Execute with helper-managed optional RLM routing.
      import inspect
      if inspect.iscoroutinefunction(self._run_with_optional_rlm_async):
        try:
          loop = asyncio.get_event_loop()
          if loop.is_running():
            # If loop is running, create a task
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as executor:
              future = executor.submit(asyncio.run, self._run_with_optional_rlm_async(user_id, session_id, user_message))
              response_text = future.result()
          else:
            response_text = loop.run_until_complete(self._run_with_optional_rlm_async(user_id, session_id, user_message))
        except RuntimeError:
          # No event loop, create one
          response_text = asyncio.run(self._run_with_optional_rlm_async(user_id, session_id, user_message))
      else:
        response_text = self._run_with_optional_rlm_async(user_id, session_id, user_message)
      
{% if spec.output_fields %}
{% set first_output = spec.output_fields[0] %}
      return {"{{ first_output.name }}": response_text}
{% else %}
      return {"response": response_text}
{% endif %}
    
    except Exception as e:
      import traceback
      error_msg = f"Error executing Google ADK agent: {str(e)}"
      print(f"⚠️  {error_msg}")
      traceback.print_exc()
{% if spec.output_fields %}
{% set first_output = spec.output_fields[0] %}
      return {"{{ first_output.name }}": f"Error: {str(e)}"}
{% else %}
      return {"response": f"Error: {str(e)}"}
{% endif %}
  
  async def _run_agent_async(self, user_id: str, session_id: str, user_message: str) -> str:
    """
    Execute agent asynchronously and collect response.
    
    Args:
      user_id: User ID
      session_id: Session ID
      user_message: User's message
    
    Returns:
      Agent's response text
    """
    # Create session
    session = await self._runner.session_service.create_session(
      app_name=f"superoptix_{{ agent_name }}",
      user_id=user_id,
    )
    
    # Create message content
    content = types.Content(
      role='user',
      parts=[types.Part.from_text(text=user_message)]
    )
    
    # Run agent and collect responses
    response_parts = []
    async for event in self._runner.run_async(
      user_id=user_id,
      session_id=session.id,
      new_message=content,
    ):
      if event.content and event.content.parts:
        for part in event.content.parts:
          if part.text and event.author == self._agent.name:
            response_parts.append(part.text)
    
    return ''.join(response_parts) if response_parts else "No response generated"
  
  async def _run_with_optional_rlm_async(self, user_id: str, session_id: str, user_message: str) -> str:
    """
    Execute with optional RLM wrapper.
    Keeps optimized pipeline aligned with minimal-mode RLM behavior.
    """
    del session_id  # Session IDs are managed by runner/session service
    # Convert config to dict if needed
    import types as python_types
    config = vars(self.config) if isinstance(self.config, python_types.SimpleNamespace) else self.config
    if not isinstance(config, dict):
      config = {}
    model_str = str(config.get("model", "gemini-2.0-flash"))
    app_name = f"superoptix_{{ agent_name }}"
    return await superoptix_run_agent_with_optional_rlm(
      agent=self._agent,
      runner=self._runner,
      prompt=user_message,
      spec_data=FULL_SPEC,
      model_name=model_str,
      app_name=app_name,
      user_id=user_id,
      logfire_enabled=bool((FULL_SPEC.get("logfire", {}) or {}).get("enabled", True)),
    )
  
  def update(self, new_variable: Any) -> None:
    """
    Update instruction (called by GEPA optimizer during optimization).
    
    Args:
      new_variable: New instruction text
    """
    super().update(new_variable)
    # Force re-initialization with new instruction
    self._agent = None
    self._runner = None
  
  def __repr__(self) -> str:
    return (
      f"{{ agent_name | to_pascal_case }}Component("
      f"framework=google_adk, "
      f"optimizable={self.optimizable})"
    )


# ======================================================================
# Factory Function
# ======================================================================

def create_{{ agent_name }}_agent(
  instruction: Optional[str] = None,
  model_config: Optional[Dict] = None,
  **kwargs
) -> {{ agent_name | to_pascal_case }}Component:
  """
  Factory function to create {{ agent_name }} agent.
  
  Args:
    instruction: Optional custom instruction (overrides default)
    model_config: Model configuration
    **kwargs: Additional configuration
  
  Returns:
    Initialized {{ agent_name | to_pascal_case }}Component
  """
  tools = []  # Tools can be added here
  
  return {{ agent_name | to_pascal_case }}Component(
    instruction=instruction,
    model_config=model_config,
    tools=tools,
    **kwargs
  )


# ======================================================================
# {{ agent_name | to_pascal_case }}Pipeline - SuperOptiX Workflow Support
# ======================================================================

import yaml
from pathlib import Path


class {{ agent_name | to_pascal_case }}Pipeline:
  """
  Google ADK pipeline with full SuperOptiX workflow support.
  
  Supports:
  - compile: Generate this code from playbook
  - evaluate: Run BDD scenarios
  - optimize: GEPA optimization  
  - run: Execute agent
  
  This makes Google ADK work with standard SuperOptiX commands!
  """
  
  def __init__(self, playbook_path: str = None):
    """Initialize pipeline from playbook."""
    # Load playbook
    if playbook_path:
      with open(playbook_path) as f:
        playbook = yaml.safe_load(f)
        self.spec = playbook.get("spec", {})
        self.metadata = playbook.get("metadata", {})
    else:
      self.spec = {}
      self.metadata = {}
    
    # Get model config from spec
    model_config = {}
    if "language_model" in self.spec:
      lm = self.spec["language_model"]
      model_config = {
        "model": lm.get("model", "gemini-2.0-flash"),
        "description": self.metadata.get("description", "AI Agent"),
      }
    
    # Create component
    self.component = create_{{ agent_name }}_agent(model_config=model_config)
    
    # Load BDD scenarios
    self.test_scenarios = self._load_bdd_scenarios()
    
    # Alias for CLI compatibility
    self.test_examples = self.test_scenarios
    
    # Metrics
    self.is_trained = False
  
  def _load_bdd_scenarios(self) -> List[Dict]:
    """Load BDD test scenarios from playbook."""
    scenarios = []
    feature_specs = self.spec.get("feature_specifications", {})
    scenario_list = feature_specs.get("scenarios", [])
    
    for scenario in scenario_list:
      scenarios.append({
        "name": scenario.get("name", "Unnamed"),
        "description": scenario.get("description", ""),
        "input": scenario.get("input", {}),
        "expected_output": scenario.get("expected_output", {}),
      })
    
    return scenarios
  
  def run(self, **inputs) -> Dict[str, Any]:
    """
    Execute agent (unified interface for SuperOptiX CLI).
    
    Args:
      **inputs: Input fields (e.g., query="...", message="...")
    
    Returns:
      Dict with output fields
    """
    # Use component's forward method
    result = self.component.forward(**inputs)
    return result
  
  def evaluate(self) -> Dict[str, Any]:
    """
    Evaluate agent with BDD scenarios.
    
    Returns:
      Dict with evaluation results
    """
    if not self.test_scenarios:
      print("⚠️  No BDD scenarios found in playbook")
      return {"pass": 0, "fail": 0, "pass_rate": 0.0}
    
    print(f"\n🔍 Evaluating {{ metadata.name }}...\n")
    print(f"Testing {len(self.test_scenarios)} BDD scenarios:\n")
    
    passed = 0
    failed = 0
    results = []
    
    for scenario in self.test_scenarios:
      # Convert SimpleNamespace to dict if needed
      import types as python_types
      if isinstance(scenario, python_types.SimpleNamespace):
        scenario = vars(scenario)
      
      scenario_name = scenario.get("name", "Unnamed") if isinstance(scenario, dict) else "Unnamed"
      
      try:
        # Get inputs and expected outputs
        inputs = scenario.get("input", {}) if isinstance(scenario, dict) else {}
        expected = scenario.get("expected_output", {}) if isinstance(scenario, dict) else {}
        
        # Run scenario
        result = self.run(**inputs)
        
        # Check if output matches expected
        success = self._evaluate_output(result, expected)
        
        if success:
          print(f"✅ {scenario_name}: PASS")
          passed += 1
        else:
          print(f"❌ {scenario_name}: FAIL")
          failed += 1
        
        results.append({
          "scenario": scenario_name,
          "passed": success,
          "output": result,
        })
      
      except Exception as e:
        import traceback
        print(f"❌ {scenario_name}: ERROR - {e}")
        traceback.print_exc()
        failed += 1
        results.append({
          "scenario": scenario_name,
          "passed": False,
          "error": str(e),
        })
    
    pass_rate = passed / len(self.test_scenarios) if self.test_scenarios else 0.0
    
    print(f"\n{'='*60}")
    print(f"Overall: {passed}/{len(self.test_scenarios)} PASS ({pass_rate*100:.1f}%)")
    print(f"{'='*60}\n")
    
    return {
      "pass": passed,
      "fail": failed,
      "total": len(self.test_scenarios),
      "pass_rate": pass_rate,
      "results": results,
    }
  
  def _evaluate_output(self, result: Dict, expected: Dict) -> bool:
    """Check if result matches expected output."""
    # Convert SimpleNamespace to dict if needed
    import types as python_types
    if isinstance(result, python_types.SimpleNamespace):
      result = vars(result)
    if isinstance(expected, python_types.SimpleNamespace):
      expected = vars(expected)
    
    # Ensure we have dicts
    if not isinstance(result, dict):
      result = {"response": str(result)}
    if not isinstance(expected, dict):
      expected = {}
    
    # Simple keyword matching
    result_str = str(result).lower()
    expected_keywords = expected.get("expected_keywords", [])
    
    if expected_keywords:
      # Convert all keywords to strings and lowercase
      keywords_str = [str(kw).lower() for kw in expected_keywords if kw]
      matches = sum(1 for kw in keywords_str if kw in result_str)
      return matches >= len(keywords_str) * 0.5  # 50% threshold
    
    # Check response field
    if "response" in expected and "response" in result:
      return expected["response"].lower() in result["response"].lower()
    
    return True  # Pass if no clear criteria
  
  def run_bdd_test_suite(
    self, auto_tune: bool = False, ignore_checks: bool = False
  ) -> Dict[str, Any]:
    """
    Run BDD test suite (CLI compatibility method).
    
    This method is called by the CLI evaluate command and wraps the 
    standard evaluate() method to format results for CLI display.
    
    Args:
      auto_tune: Whether to automatically optimize based on results (unused)
      ignore_checks: If True, treats all scenarios as passed
    
    Returns:
      Dict containing test results formatted for CLI
    """
    if not self.test_examples:
      return {
        "success": False,
        "message": "No BDD specifications defined in feature_specifications",
        "summary": {"total": 0, "passed": 0, "failed": 0, "pass_rate": "0.00%"},
        "bdd_results": {"detailed_results": []},
        "model_analysis": {},
        "recommendations": [],
      }
    
    # Run evaluation
    eval_results = self.evaluate()
    
    # Convert to CLI format
    detailed_results = []
    for result in eval_results.get("results", []):
      passed = result.get("passed", False)
      if ignore_checks:
        passed = True
      
      detailed_results.append({
        "scenario_name": result.get("scenario", "Unnamed"),
        "description": result.get("description", ""),
        "passed": passed,
        "confidence_score": 1.0 if passed else 0.0,
        "actual_output": result.get("output", {}),
      })
    
    # Calculate metrics
    total = eval_results.get("total", 0)
    passed = eval_results.get("pass", 0) if not ignore_checks else total
    failed = total - passed
    pass_rate = eval_results.get("pass_rate", 0.0) * 100
    
    return {
      "success": True,
      "summary": {
        "total": total,
        "passed": passed,
        "failed": failed,
        "pass_rate": f"{pass_rate:.1f}%",
      },
      "bdd_results": {
        "detailed_results": detailed_results,
        "total_scenarios": total,
        "scenarios_passed": passed,
        "scenarios_failed": failed,
        "pass_rate": f"{pass_rate:.1f}%",
        "bdd_score": pass_rate / 100,
      },
      "model_analysis": {
        "framework": "Google ADK",
        "model": "{{ spec.language_model.model | default('gemini-2.0-flash') }}",
      },
      "recommendations": [],
    }
  
  def optimize_with_gepa(
    self,
    auto: str = "medium",
    metric: str = "response_accuracy",
  ) -> Dict[str, Any]:
    """
    Optimize agent with GEPA.
    
    Args:
      auto: Optimization budget ("light", "medium", "heavy")
      metric: Metric to optimize
    
    Returns:
      Optimization results
    """
    from superoptix.optimizers.universal_gepa import UniversalGEPA
    
    print(f"\n🔄 Optimizing {{ metadata.name }} with GEPA...\n")
    print(f"Framework: Google ADK")
    print(f"Optimization level: {auto}")
    print(f"Optimizing: instruction\n")
    
    # Define metric function
    def eval_metric(inputs, outputs, gold, component_name=None):
      """Metric function for GEPA."""
      # Simple keyword-based evaluation
      output_str = str(outputs).lower()
      expected_keywords = gold.get("expected_keywords", [])
      
      if expected_keywords:
        matches = sum(1 for kw in expected_keywords if str(kw).lower() in output_str)
        score = matches / len(expected_keywords)
      else:
        score = 0.5  # Default
      
      return score
    
    # Prepare training data from BDD scenarios
    trainset = []
    for scenario in self.test_scenarios[:len(self.test_scenarios)//2]:  # First half for training
      trainset.append({
        "inputs": scenario["input"],
        "expected": scenario["expected_output"],
      })
    
    # Validation set
    valset = []
    for scenario in self.test_scenarios[len(self.test_scenarios)//2:]:  # Second half for validation
      valset.append({
        "inputs": scenario["input"],
        "expected": scenario["expected_output"],
      })
    
    if not trainset:
      print("⚠️  No training data available. Need BDD scenarios in playbook.")
      return {"error": "No training data"}
    
    # Run GEPA optimization
    try:
      optimizer = UniversalGEPA(
        metric=eval_metric,
        auto=auto,
        reflection_lm="gpt-4o-mini",
      )
      
      result = optimizer.compile(
        component=self.component,
        trainset=trainset,
        valset=valset if valset else trainset,
      )
      
      print(f"\n✅ Optimization complete!")
      print(f"Best score: {result.best_score:.2f}")
      print(f"Optimized instruction:\n{result.optimized_component.variable[:300]}...\n")
      
      # Update component with optimized instruction
      self.component = result.optimized_component
      self.is_trained = True
      
      return {
        "best_score": result.best_score,
        "optimized_instruction": result.optimized_component.variable,
        "improvement": result.best_score - result.initial_score if hasattr(result, 'initial_score') else 0.0,
      }
    
    except Exception as e:
      print(f"❌ Optimization failed: {e}")
      import traceback
      traceback.print_exc()
      return {"error": str(e)}


# ======================================================================
# Main Execution
# ======================================================================

if __name__ == "__main__":
  """
  Test the pipeline directly.
  
  Usage:
    python {{ agent_name }}_pipeline.py
  """
  print("🧪 Testing {{ metadata.name }} Pipeline\n")
  
  # Find playbook
  playbook_path = Path(__file__).parent.parent / "playbook" / "{{ agent_name }}_playbook.yaml"
  
  if not playbook_path.exists():
    print(f"❌ Playbook not found: {playbook_path}")
    print("💡 Run from agent directory or compile with: super agent compile {{ agent_name }}")
    exit(1)
  
  # Initialize pipeline
  pipeline = {{ agent_name | to_pascal_case }}Pipeline(playbook_path=str(playbook_path))
  
  # Test run
  print("🚀 Testing Google ADK agent execution...\n")
{% if spec.input_fields %}
{% set first_input = spec.input_fields[0] %}
  test_result = pipeline.run({{ first_input.name }}="What is AI?")
{% else %}
  test_result = pipeline.run(query="What is AI?")
{% endif %}
  print(f"Result: {test_result}\n")
  
  # Evaluate if scenarios available
  if pipeline.test_scenarios:
    print("🧪 Running evaluation...\n")
    eval_result = pipeline.evaluate()
    print(f"\nEvaluation: {eval_result['pass']}/{eval_result['total']} passed")
