"""
Flask web application for the OpenRouter Agent.
"""

import os
import uuid
import json
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from flask import Flask, jsonify, request, render_template

from core.agent import OpenRouterAgent
from core.models import Message, Role
from core.tools.builtin import get_builtin_tools, set_working_directory
from core.tools.custom import get_custom_tools


app = Flask(__name__)
app.secret_key = os.getenv("SECRET_KEY", str(uuid.uuid4()))

PROMPTS_FILE = os.path.join(os.path.dirname(__file__), "prompts.json")


def load_prompts() -> Dict[str, str]:
    """Load saved prompts from file"""
    if os.path.exists(PROMPTS_FILE):
        try:
            with open(PROMPTS_FILE, "r") as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            return {}
    return {}


def save_prompts(prompts: Dict[str, str]) -> None:
    """Save prompts to file"""
    with open(PROMPTS_FILE, "w") as f:
        json.dump(prompts, f, indent=2)


@dataclass
class AgentSession:
    """Stores agent state for a session"""
    agent: OpenRouterAgent
    conversation_history: List[Message] = field(default_factory=list)
    pending_question: Optional[str] = None


sessions: Dict[str, AgentSession] = {}


def create_agent(working_directory: Optional[str] = None) -> OpenRouterAgent:
    """Create an agent with all tools registered"""
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        raise ValueError("OPENROUTER_API_KEY not set")
    
    if working_directory:
        set_working_directory(working_directory)
    
    tools = get_builtin_tools()
    tools.update(get_custom_tools())
    
    return OpenRouterAgent(
        api_key=api_key,
        model=os.getenv("OPENROUTER_MODEL", "qwen/qwen3-coder-next"),
        temperature=0.7,
        max_iterations=8,
        tools=tools,
    )


@app.route("/")
def index():
    """Serve the main page"""
    return render_template("index.html")


@app.route("/api/chat", methods=["POST"])
def chat():
    """Handle chat messages"""
    data = request.get_json()
    user_message = data.get("message", "")
    session_id = data.get("session_id")
    working_directory = data.get("working_directory")
    
    if not session_id or session_id not in sessions:
        session_id = str(uuid.uuid4())
        agent = create_agent(working_directory)
        sessions[session_id] = AgentSession(agent=agent)
    else:
        agent = sessions[session_id].agent
    
    agent.conversation_history = sessions[session_id].conversation_history
    
    if sessions[session_id].pending_question:
        return jsonify({
            "type": "error",
            "message": "Please answer the pending question first"
        })
    
    session = sessions[session_id]
    
    system_prompt = """You are a helpful AI assistant with access to various tools.
    When asked a question, use the available tools to gather information or perform actions.
    Always be concise and helpful in your responses."""
    
    result = run_agent_with_handling(agent, user_message, system_prompt, session)
    
    session.conversation_history = agent.conversation_history
    
    return jsonify({
        "session_id": session_id,
        **result
    })


@app.route("/api/chat/respond", methods=["POST"])
def chat_respond():
    """Respond to a pending question"""
    data = request.get_json()
    session_id = data.get("session_id")
    response = data.get("response", "")
    
    if not session_id or session_id not in sessions:
        return jsonify({
            "type": "error",
            "message": "Invalid session"
        })
    
    session = sessions[session_id]
    
    if not session.pending_question:
        return jsonify({
            "type": "error",
            "message": "No pending question"
        })
    
    question = session.pending_question
    session.pending_question = None
    
    user_message = f"User's answer to '{question}': {response}"
    
    system_prompt = """You are a helpful AI assistant with access to various tools.
    When asked a question, use the available tools to gather information or perform actions.
    Always be concise and helpful in your responses."""
    
    result = run_agent_with_handling(session.agent, user_message, system_prompt, session)
    
    session.conversation_history = session.agent.conversation_history
    
    return jsonify({
        "session_id": session_id,
        **result
    })


def run_agent_with_handling(
    agent: OpenRouterAgent,
    user_input: str,
    system_prompt: str,
    session: AgentSession
) -> Dict[str, Any]:
    """Run the agent, handling ask_user tool specially"""
    
    if system_prompt and not any(m.role == Role.SYSTEM for m in agent.conversation_history):
        agent.add_message(Message(role=Role.SYSTEM, content=system_prompt))
    
    agent.add_message(Message(role=Role.USER, content=user_input))
    
    iteration = 0
    final_response = None
    
    while iteration < agent.max_iterations:
        iteration += 1
        
        messages = [msg.to_dict() for msg in agent.conversation_history]
        
        tools_list = None
        if agent.tools:
            tools_list = [tool.to_dict() for tool in agent.tools.values()]
        
        response = agent._make_api_call(messages, tools_list)
        
        choice = response["choices"][0]
        message = choice["message"]
        
        if "tool_calls" in message and message["tool_calls"]:
            assistant_message = Message(
                role=Role.ASSISTANT,
                content=message.get("content"),
                tool_calls=message["tool_calls"]
            )
            
            agent.add_message(assistant_message)
            
            for tool_call in message["tool_calls"]:
                func_name = tool_call["function"]["name"]
                func_args_raw = tool_call["function"]["arguments"]
                if isinstance(func_args_raw, str):
                    import json
                    func_args = json.loads(func_args_raw)
                else:
                    func_args = func_args_raw
                
                if func_name == "ask_user":
                    question = func_args.get("question", "")
                    session.pending_question = question
                    
                    return {
                        "type": "need_input",
                        "question": question
                    }
                
                if func_name in agent.tools:
                    func = agent.tools[func_name].function
                    try:
                        result = func(**func_args)
                    except Exception as e:
                        result = f"Error: {str(e)}"
                else:
                    result = f"Error: Unknown tool {func_name}"
                
                tool_message = Message(
                    role=Role.TOOL,
                    content=str(result),
                    tool_call_id=tool_call["id"]
                )
                agent.add_message(tool_message)
            
            continue
        else:
            assistant_message = Message(
                role=Role.ASSISTANT,
                content=message.get("content", "")
            )
            
            agent.add_message(assistant_message)
            final_response = message.get("content", "")
            break
    
    if final_response is None:
        final_response = "Maximum iterations reached without final response."
    
    return {
        "type": "response",
        "message": final_response
    }


@app.route("/api/chat/history", methods=["GET"])
def chat_history():
    """Get conversation history"""
    session_id = request.args.get("session_id")
    
    if not session_id or session_id not in sessions:
        return jsonify({"history": []})
    
    history = []
    for msg in sessions[session_id].agent.conversation_history:
        history.append({
            "role": msg.role.value if hasattr(msg.role, 'value') else msg.role,
            "content": msg.content
        })
    
    return jsonify({"history": history})


@app.route("/api/chat/clear", methods=["POST"])
def chat_clear():
    """Clear a session"""
    data = request.get_json()
    session_id = data.get("session_id")
    
    if session_id and session_id in sessions:
        del sessions[session_id]
    
    return jsonify({"status": "ok"})


@app.route("/api/prompts", methods=["GET"])
def get_prompts():
    """Get all saved prompts"""
    prompts = load_prompts()
    return jsonify({"prompts": prompts})


@app.route("/api/prompts", methods=["POST"])
def save_prompt():
    """Save a new prompt"""
    data = request.get_json()
    name = data.get("name", "").strip()
    prompt = data.get("prompt", "").strip()
    
    if not name:
        return jsonify({"error": "Prompt name is required"}), 400
    if not prompt:
        return jsonify({"error": "Prompt content is required"}), 400
    
    prompts = load_prompts()
    prompts[name] = prompt
    save_prompts(prompts)
    
    return jsonify({"status": "ok", "name": name})


@app.route("/api/prompts/<prompt_name>", methods=["DELETE"])
def delete_prompt(prompt_name):
    """Delete a saved prompt"""
    prompts = load_prompts()
    if prompt_name in prompts:
        del prompts[prompt_name]
        save_prompts(prompts)
        return jsonify({"status": "ok"})
    return jsonify({"error": "Prompt not found"}), 404


if __name__ == "__main__":
    port = int(os.getenv("PORT", 5000))
    app.run(host="0.0.0.0", port=port, debug=True)