"""
Project Name: Autonomous AI Orchestrator

Description:
An end-to-end, self-improving AI agent framework that takes user goals and autonomously:

- Plans multi-step workflows (Research, Learn, Develop, Execute, Diagnose, Review)
- Performs web searches and aggregates reliable information
- Generates and safely executes Python code
- Self-evaluates and optimizes its own processes
- Provides a Learning Coach for personalized, source-backed teaching
- Persists session memory, logs activities, and snapshots full state for audit and rollback

Instructions:
1. **Dependencies**: Ensure your `requirements.txt` contains:
   ```
   duckduckgo-search
   gradio
   requests
   python-dotenv
   feedparser
   ```
2. **Module Structure**: Place the following modules alongside your main script (`app.py`):
   - `memory.py`
   - `planner.py`
   - `executor.py`
   - `critic.py`
   - `cognitive_engine.py`
   - `web_searcher.py`
   - `hf_packager.py`
   - (Optional CSS: `style.css` for UI styling)
3. **Run the Application**:
   ```bash
   python app.py
   ```
4. **Access Interfaces**:
   - **Gradio UI**: Navigate to `http://<host>:<port>/ui` for Task Execution and Learning Coach tabs
   - **API Endpoints**:
     - `GET /status` — Health check and available routes
     - `GET /generate?prompt=<text>` — Quick text generation
     - `POST /generate` — JSON-based generation with session tracking
     - `GET /memory/{session_id}` — Retrieve session history
     - `GET /summary/{session_id}` — Summarize recent interactions
5. **Persistence & Logs**:
   - Activity logs in `log.txt`
   - Session memory stored in `memory.json`
   - Snapshots saved under `snapshots/` with recoverable JSON files
"""

import os
import zipfile
import base64
import gradio as gr
import uuid
import json
import logging
import datetime
from memory import MemoryManager
from planner import Planner
from executor import Executor
from critic import Critic
from cognitive_engine import CognitiveEngine
from web_searcher import WebSearcher
from hf_packager import HFSpacePackager

# Initialize components
memory = MemoryManager()
planner = Planner()
executor = Executor()
critic = Critic()
cog_engine = CognitiveEngine()
web_searcher = WebSearcher()
packager = HFSpacePackager()

# Set up logging
logging.basicConfig(filename='log.txt', level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')

class AutonomousAgent:
    def __init__(self):
        self.state_file = "state.json"
        self.load_state()
        
    def load_state(self):
        try:
            with open(self.state_file, 'r') as f:
                self.state = json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            self.state = {"sessions": {}}
            self.save_state()
    
    def save_state(self):
        with open(self.state_file, 'w') as f:
            json.dump(self.state, f, indent=2)
    
    def process_goal(self, goal, session_id=None):
        try:
            if not session_id:
                session_id = str(uuid.uuid4())
                self.state["sessions"][session_id] = {"goal": goal, "status": "processing"}
                memory.init_session(session_id)
                self.save_state()
            
            # Add to memory
            memory.add(session_id, "user_goal", goal)
            
            # Plan the task
            plan = planner.plan_task(goal, memory.get(session_id))
            memory.add(session_id, "plan", plan)
            
            # Execute plan
            results = []
            for step in plan:
                key = step.split(":")[0].lower()
                if key in ["research", "search"]:
                    query = step.split(":",1)[1].strip()
                    res = web_searcher.search(query)
                    memory.add(session_id, f"search:{query}", res)
                    results.append(f"🔍 Search '{query}': {res[:300]}...")
                elif key in ["develop", "code"]:
                    code = cog_engine.generate_code(step, memory.get(session_id))
                    out = executor.execute_code(code)
                    memory.add(session_id, "generated_code", code)
                    memory.add(session_id, "execution_result", out)
                    review = critic.review(step, out)
                    memory.add(session_id, "review", review)
                    if "error" in review.lower():
                        improved = cog_engine.improve_code(code, review)
                        memory.add(session_id, "enhanced_code", improved)
                        out = executor.execute_code(improved)
                        results.append(f"🛠️ Enhanced code: {out}")
                    else:
                        results.append(f"✅ Code executed: {out}")
                elif key in ["diagnose", "check"]:
                    issues = cog_engine.identify_improvements(step)
                    memory.add(session_id, "diagnosis", issues)
                    if issues:
                        fixes = cog_engine.generate_enhancements(issues)
                        cog_engine.apply_enhancements(fixes)
                        results.append(f"⚙️ Repaired: {', '.join(issues)}")
                    else:
                        results.append("✅ Healthy")
            
            self.state["sessions"][session_id]["status"] = "completed"
            self.save_state()
            snap = packager.create_snapshot({"session_id": session_id, "memory": memory.get(session_id), "results": results})
            return "\n\n".join(results), session_id, snap
        except Exception as e:
            logging.error(f"Error: {e}", exc_info=True)
            issues = [str(e)]
            fixes = cog_engine.generate_enhancements(issues)
            cog_engine.apply_enhancements(fixes)
            return f"⚠️ Error occurred: {e}", session_id, ""

agent = AutonomousAgent()

with gr.Blocks(css="style.css", title="Autonomous AI") as demo:
    sid = gr.State()
    hist = gr.State({"history": []})
    gr.Markdown("# 🤖 Autonomous AI System")
    with gr.Tab("Task Execution"):
        inp = gr.Textbox(label="Your Goal")
        btn = gr.Button("Execute Goal")
        out = gr.Textbox(label="Results", lines=10)
        sid_box = gr.Textbox(label="Session ID")
        snap_box = gr.Textbox(label="Snapshot URL")
        btn.click(agent.process_goal, [inp, sid], [out, sid_box, snap_box])
    with gr.Tab("Learning Coach"):
        chat = gr.Chatbot()
        q_in = gr.Textbox(label="Your Question")
        ask = gr.Button("Ask")
        clear = gr.Button("Clear")
        kb = gr.Markdown("")
        ask.click(lambda q,h: (h+[(q, cog_engine.learning_coach_response(q))], h), [q_in, hist], [chat, hist])
        clear.click(lambda h: ([],h), [hist], [chat, hist])

demo.launch()