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

from abi_core.common.utils import yield_chunk_data

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
                    ):
                        async for sse_bytes in yield_chunk_data(chunk):
                            yield sse_bytes
                    
                    # 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