#!/usr/bin/env python3
"""
Antigravity Workflow Engine
===========================
The core execution engine for autonomous genesis system development.
Synthesizes high-level goals into executable steps.

Architecture:
- Plan: Synthesized by Gemini (Pro/Flash)
- Execute: Step-by-step implementation with verification
- Memory: Recursive integration of results into the knowledge base
"""

import json
import logging
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, asdict

# Path setup
GENESIS_ROOT = Path("/mnt/e/genesis-system")
sys.path.insert(0, str(GENESIS_ROOT))

from core.gemini_executor import GeminiExecutor

logger = logging.getLogger("core.antigravity_workflow")

@dataclass
class WorkflowStep:
    id: str
    description: str
    action: str
    status: str = "pending"
    result: Optional[str] = None
    error: Optional[str] = None

@dataclass
class WorkflowPlan:
    plan_id: str
    goal: str
    steps: List[WorkflowStep]
    created_at: str
    status: str = "planned"

class AntigravityWorkflow:
    """
    Orchestrates goal-to-execution cycles for Genesis.
    """
    
    def __init__(self):
        self.executor = GeminiExecutor(use_rate_maximizer=True)
        self.plans_dir = GENESIS_ROOT / "data" / "antigravity-sync" / "plans"
        self.exec_dir = GENESIS_ROOT / "data" / "antigravity-sync" / "executions"
        self.plans_dir.mkdir(parents=True, exist_ok=True)
        self.exec_dir.mkdir(parents=True, exist_ok=True)

    def create_plan(self, goal: str, context: dict = None) -> WorkflowPlan:
        """
        Synthesize a goal into a multi-step plan using Gemini.
        """
        prompt = f"""# GOAL: {goal}
# CONTEXT: {json.dumps(context or {})}

Act as the Genesis System Architect. Break this goal down into exact, executable technical steps.
Return ONLY a JSON list of steps. Each step MUST have:
- "id": unique string ID
- "description": clear explanation of what to do
- "action": the specific technical implementation instruction

Response Format:
[
  {{"id": "step_1", "description": "...", "action": "..."}},
  ...
]
"""
        response = self.executor.execute_optimized(prompt, task_type="architecture")
        
        if not response.success:
            raise RuntimeError(f"Failed to create plan: {response.error}")
            
        try:
            # Extract JSON from response (handling potential markdown formatting)
            cleaned_resp = response.response.strip()
            if "```json" in cleaned_resp:
                cleaned_resp = cleaned_resp.split("```json")[1].split("```")[0].strip()
            elif "```" in cleaned_resp:
                cleaned_resp = cleaned_resp.split("```")[1].split("```")[0].strip()
                
            steps_data = json.loads(cleaned_resp)
            steps = [WorkflowStep(**s) for s in steps_data]
            
            plan_id = f"plan_{int(time.time())}"
            plan = WorkflowPlan(
                plan_id=plan_id,
                goal=goal,
                steps=steps,
                created_at=datetime.now(timezone.utc).isoformat()
            )
            return plan
        except Exception as e:
            logger.error(f"Error parsing plan: {e}")
            logger.error(f"Raw response: {response.response}")
            raise RuntimeError(f"Plan synthesis failed: {e}")

    def execute_plan(self, plan_data: dict) -> dict:
        """
        Execute a synthesized plan.
        """
        plan_id = plan_data.get("plan_id")
        steps = plan_data.get("steps", [])
        results = []
        
        print(f"[Workflow] Executing plan: {plan_id}")
        
        for step_data in steps:
            step = WorkflowStep(**step_data)
            print(f"  > Step {step.id}: {step.description}")
            
            # Execute step via Gemini
            exec_prompt = f"""## GOAL: {plan_data.get('goal')}
## STEP DESCRIPTION: {step.description}
## ACTION: {step.action}

Implement this step. If it involves code, provide the code. 
Verify the implementation.
If complete, end with TASK_COMPLETE.
"""
            resp = self.executor.execute_optimized(exec_prompt, task_type="code_generation")
            
            if resp.success:
                step.status = "completed"
                step.result = resp.response
            else:
                step.status = "failed"
                step.error = resp.error
            
            results.append(asdict(step))
            
            if step.status == "failed":
                print(f"  X Step {step.id} FAILED: {step.error}")
                break
                
        return {
            "plan_id": plan_id,
            "status": "completed" if all(r['status'] == "completed" for r in results) else "partially_failed",
            "results": results
        }

if __name__ == "__main__":
    # Test plan creation
    workflow = AntigravityWorkflow()
    # plan = workflow.create_plan("Initialize the agency outreach scraper for 3 targets.")
    # print(json.dumps(asdict(plan), indent=2))