"""
Web Interface for ABI Agents - Compatible with Open WebUI
Generated by ABI-Core scaffolding v1.2.0
"""

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from typing import Optional, List
import asyncio
import json
import time
import logging

logger = logging.getLogger(__name__)

class ChatMessage(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    stream: Optional[bool] = False

class AgentWebInterface:
    """
    Generic web interface for ABI agents
    
    Provides FastAPI endpoints compatible with:
    - Open WebUI (Ollama-compatible API)
    - Server-Sent Events (SSE) streaming
    - Standard REST API
    """
    
    def __init__(self, agent_instance, interface_name: str = "Agent Web Interface"):
        self.agent_instance = agent_instance
        self.interface_name = interface_name
        self.agent_model_name = f"{interface_name.lower().replace(' ', '-')}:latest"
        
        self.app = FastAPI(
            title=f"{interface_name}",
            description=f"Web interface for {interface_name} - Compatible with Open WebUI",
            version="1.4.0"
        )
        self.setup_routes()
        
        logger.info(f"🌐 {interface_name} web interface initialized (Open WebUI compatible)")

    def setup_routes(self):
        """Setup FastAPI routes for the web interface"""
        
        # ============ ENDPOINTS PARA OPEN WEBUI ============
        @self.app.get("/api/tags")
        async def list_models():
            """Lista modelos disponibles (compatible con Open WebUI)"""
            return {
                "models": [{
                    "name": self.agent_model_name,
                    "model": self.agent_model_name,
                    "modified_at": "2024-12-15T00:00:00Z",
                    "size": 0,
                    "digest": f"sha256:{self.interface_name.lower().replace(' ', '')}",
                    "details": {
                        "parent_model": "",
                        "format": "abi-agent",
                        "family": "abi",
                        "families": ["abi"],
                        "parameter_size": "agent-based",
                        "quantization_level": "native"
                    },
                    "version": "1.4.0"
                }]
            }

        @self.app.get("/api/version")
        async def get_version():
            """Retorna versión compatible con Ollama"""
            return {"version": "1.4.0"}

        @self.app.get("/api/ps")
        async def list_running():
            """Lista modelos en ejecución (compatible con Open WebUI)"""
            return {
                "models": [{
                    "name": self.agent_model_name,
                    "model": self.agent_model_name,
                    "size": 0,
                    "digest": f"sha256:{self.interface_name.lower().replace(' ', '')}",
                    "expires_at": "2099-12-31T23:59:59Z"
                }]
            }

        @self.app.post("/api/chat")
        async def chat(request: ChatRequest):
            """Chat endpoint compatible con Open WebUI"""
            # Extraer el último mensaje del usuario
            user_message = request.messages[-1].content if request.messages else ""
            
            if request.stream:
                return StreamingResponse(
                    self._stream_response_openwebui(user_message),
                    media_type="text/plain",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "close",
                        "Access-Control-Allow-Origin": "*",
                        "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
                        "Access-Control-Allow-Headers": "Content-Type",
                    }
                )
            else:
                result = await self._process_query_v2(user_message)
                return {
                    "model": self.agent_model_name,
                    "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    "message": {
                        "role": "assistant",
                        "content": result
                    },
                    "done": True
                }

        @self.app.post("/api/generate")
        async def generate(request: dict):
            """Generate endpoint compatible con Open WebUI"""
            prompt = request.get("prompt", "")
            stream = request.get("stream", False)
            
            if stream:
                return StreamingResponse(
                    self._stream_response_openwebui(prompt),
                    media_type="text/plain",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "close",
                        "Access-Control-Allow-Origin": "*",
                        "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
                        "Access-Control-Allow-Headers": "Content-Type",
                    }
                )
            else:
                result = await self._process_query_v2(prompt)
                return {
                    "model": self.agent_model_name,
                    "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    "response": result,
                    "done": True
                }

        # ============ ENDPOINTS ORIGINALES ============
        @self.app.get("/")
        async def root():
            """Root endpoint with interface information"""
            return {
                "interface": self.interface_name,
                "status": "running",
                "model": self.agent_model_name,
                "version": "1.4.0",
                "compatibility": "Open WebUI",
                "endpoints": {
                    "openwebui_chat": "/api/chat - POST - Open WebUI chat endpoint",
                    "openwebui_generate": "/api/generate - POST - Open WebUI generate endpoint", 
                    "openwebui_models": "/api/tags - GET - List available models",
                    "stream": "/stream - POST - Stream agent responses (SSE)",
                    "query": "/query - POST - Single query endpoint",
                    "health": "/health - GET - Health check"
                }
            }
        
        @self.app.get("/health")
        async def health():
            """Health check endpoint"""
            return {
                "status": "healthy",
                "interface": self.interface_name,
                "model": self.agent_model_name,
                "version": "1.4.0",
                "agent_available": self.agent_instance is not None
            }

        @self.app.post("/stream")
        async def stream_query(request: dict):
            """
            Stream agent responses via Server-Sent Events (SSE)
            
            Expected request format:
            {
                "query": "Your query here",
                "context_id": "optional-context-id",
                "task_id": "optional-task-id"
            }
            """
            query = request.get("query")
            if not query:
                return {"error": "Query is required"}
            
            context_id = request.get("context_id", "web-session")
            task_id = request.get("task_id", f"task-{int(time.time())}")

            async def generate_response():
                """Generate SSE stream from agent responses"""
                try:
                    # 1) Open the SSE channel
                    yield b"event: ping\\ndata: {}\\n\\n"
                    
                    logger.info(f"🔄 Starting stream for query: {query[:100]}...")
                    
                    # 2) Stream real agent responses (MUST be async generator)
                    async for chunk in self.agent_instance.stream(
                        query=query, 
                        context_id=context_id, 
                        task_id=task_id
                    ):
                        # Convert chunk to serializable format
                        try:
                            if hasattr(chunk, 'model_dump'):
                                # Pydantic model (v2)
                                chunk_data = chunk.model_dump()
                            elif hasattr(chunk, 'dict'):
                                # Pydantic model (v1)
                                chunk_data = chunk.dict()
                            elif hasattr(chunk, '__dict__'):
                                # Regular object
                                chunk_data = chunk.__dict__
                            elif isinstance(chunk, dict):
                                # Already a dictionary
                                chunk_data = chunk
                            else:
                                # Fallback to string representation
                                chunk_data = {
                                    "message": str(chunk), 
                                    "type": type(chunk).__name__
                                }
                            
                            # Add metadata for debugging
                            chunk_data["_meta"] = {
                                "chunk_type": type(chunk).__name__,
                                "timestamp": time.time(),
                                "interface": self.interface_name
                            }
                            
                            logger.debug(f"📤 Chunk sent: {type(chunk).__name__} - {str(chunk)[:100]}...")
                            
                            yield (f"data: {json.dumps(chunk_data)}\\n\\n").encode()
                            
                        except Exception as serialize_error:
                            # If serialization fails, send error info
                            error_data = {
                                "error": "Serialization failed",
                                "chunk_type": type(chunk).__name__,
                                "details": str(serialize_error),
                                "_meta": {
                                    "timestamp": time.time(),
                                    "interface": self.interface_name
                                }
                            }
                            logger.error(f"❌ Serialization error: {serialize_error}")
                            yield (f"data: {json.dumps(error_data)}\\n\\n").encode()
                    
                    # 3) Normal completion
                    completion_data = {
                        "event": "done",
                        "message": "Stream completed successfully",
                        "_meta": {
                            "timestamp": time.time(),
                            "interface": self.interface_name
                        }
                    }
                    yield (f"event: done\\ndata: {json.dumps(completion_data)}\\n\\n").encode()
                    logger.info("✅ Stream completed successfully")
                    
                except asyncio.CancelledError:
                    # Client closed connection; exit silently
                    logger.info("🔌 Client disconnected")
                    raise
                except Exception as e:
                    # 4) Log + error event so client receives something before close
                    logger.exception("💥 Error in SSE generate_response: %s", e)
                    error_data = {
                        "error": str(e),
                        "error_type": type(e).__name__,
                        "_meta": {
                            "timestamp": time.time(),
                            "interface": self.interface_name
                        }
                    }
                    yield (f"event: error\\ndata: {json.dumps(error_data)}\\n\\n").encode()
                    # Small pause to drain buffer
                    await asyncio.sleep(0.05)

            return StreamingResponse(
                generate_response(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive",
                    "Access-Control-Allow-Origin": "*",
                    "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
                    "Access-Control-Allow-Headers": "Content-Type",
                },
            )
        
        @self.app.post("/query")
        async def single_query(request: dict):
            """
            Single query endpoint (non-streaming)
            
            Expected request format:
            {
                "query": "Your query here",
                "context_id": "optional-context-id",
                "task_id": "optional-task-id"
            }
            """
            query = request.get("query")
            if not query:
                return {"error": "Query is required"}
            
            context_id = request.get("context_id", "web-session")
            task_id = request.get("task_id", f"task-{int(time.time())}")
            
            try:
                # Collect all chunks from the stream
                chunks = []
                async for chunk in self.agent_instance.stream(
                    query=query, 
                    context_id=context_id, 
                    task_id=task_id
                ):
                    if hasattr(chunk, 'model_dump'):
                        chunks.append(chunk.model_dump())
                    elif hasattr(chunk, 'dict'):
                        chunks.append(chunk.dict())
                    elif isinstance(chunk, dict):
                        chunks.append(chunk)
                    else:
                        chunks.append({"message": str(chunk), "type": type(chunk).__name__})
                
                return {
                    "query": query,
                    "context_id": context_id,
                    "task_id": task_id,
                    "chunks": chunks,
                    "total_chunks": len(chunks),
                    "interface": self.interface_name
                }
                
            except Exception as e:
                logger.error(f"Error in single query: {e}")
                return {
                    "error": str(e),
                    "error_type": type(e).__name__,
                    "query": query,
                    "interface": self.interface_name
                }

    async def _process_query_v2(self, query: str) -> str:
        """
        Procesa query a través del agente usando solo stream()
        Compatible con v1.2.0 - Recolecta chunks del stream
        """
        try:
            # Recolectar todos los chunks del stream
            result_chunks = []
            async for chunk in self.agent_instance.stream(
                query=query, 
                context_id="openwebui-session",
                task_id=f"task-{int(time.time())}"
            ):
                # Extraer contenido del chunk
                if hasattr(chunk, 'get') and isinstance(chunk, dict):
                    content = chunk.get('content', str(chunk))
                elif hasattr(chunk, 'content'):
                    content = chunk.content
                else:
                    content = str(chunk)
                
                result_chunks.append(content)
            
            # Combinar todos los chunks en respuesta final
            if result_chunks:
                # Si el último chunk parece ser la respuesta final, usarlo
                final_result = result_chunks[-1]
                if isinstance(final_result, str) and len(final_result) > 50:
                    return final_result
                else:
                    # Combinar todos los chunks
                    return " ".join(str(chunk) for chunk in result_chunks if chunk)
            else:
                return "No response generated"
                
        except Exception as e:
            logger.error(f"Error processing query: {e}")
            return f"Error: {str(e)}"

    async def _stream_response_openwebui(self, query: str):
        """Stream response en formato compatible con Ollama/Open WebUI"""
        try:
            async for chunk in self.agent_instance.stream(
                query=query, 
                context_id="openwebui-session",
                task_id=f"task-{int(time.time())}"
            ):
                # Extraer contenido del chunk
                if hasattr(chunk, 'get') and isinstance(chunk, dict):
                    content = chunk.get('content', str(chunk))
                elif hasattr(chunk, 'content'):
                    content = chunk.content
                else:
                    content = str(chunk)

                response = {
                    "model": self.agent_model_name,
                    "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                    "message": {
                        "role": "assistant",
                        "content": content
                    },
                    "done": False
                }
                yield json.dumps(response) + "\n"

            # Mensaje final
            final = {
                "model": self.agent_model_name,
                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "message": {
                    "role": "assistant",
                    "content": ""
                },
                "done": True
            }
            yield json.dumps(final) + "\n"

        except Exception as e:
            logger.error(f"Error streaming response: {e}")
            error = {"error": str(e)}
            yield json.dumps(error) + "\n"

    def get_app(self) -> FastAPI:
        """Get the FastAPI application instance"""
        return self.app