"""
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 abi_core.session import SessionStore

from fastapi import FastAPI, Request, Header
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__)


def _extract_token(authorization: Optional[str]) -> Optional[str]:
    """Pull a bearer token out of an Authorization header, if present."""
    if not authorization:
        return None
    parts = authorization.split(None, 1)
    if len(parts) == 2 and parts[0].lower() == "bearer":
        return parts[1].strip()
    return authorization.strip() or None

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"

        # Framework-managed sessions (opt-in), backed by the SAME backend the
        # agent uses so token -> context_id and conversation context share one
        # store. With SESSION_BACKEND=redis this is LB/multi-pod safe.
        self.session_store = SessionStore(getattr(agent_instance, "session_backend", None))

        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)")

    async def _resolve_context_id(self, authorization: Optional[str]) -> str:
        """Resolve a bearer token to its backend context_id.

        The context_id is generated in the backend, never trusted from the
        client (prevents spoofing and the shared "web-session" collision). If no
        valid token is present, an anonymous session is created in the same
        backend so each caller still gets a unique, isolated context_id.
        """
        token = _extract_token(authorization)
        if token:
            session = await self.session_store.resolve(token)
            if session is not None:
                return session.context_id
        session = await self.session_store.create_session(metadata={"anonymous": True})
        return session.context_id

    def setup_routes(self):
        """Setup FastAPI routes for the web interface"""

        # ============ SESSION ENDPOINTS (framework-managed) ============
        @self.app.post("/session/start")
        async def session_start(request: dict | None = None):
            metadata = (request or {}).get("metadata") if isinstance(request, dict) else None
            session = await self.session_store.create_session(metadata=metadata or {})
            return {"session_token": session.tokens[0], "expires_at": session.expires_at}

        @self.app.post("/session/rotate")
        async def session_rotate(authorization: Optional[str] = Header(default=None)):
            token = _extract_token(authorization)
            new_token = await self.session_store.rotate(token) if token else None
            if not new_token:
                return JSONResponse(status_code=401, content={"error": "Invalid or expired session token"})
            return {"session_token": new_token}

        @self.app.post("/session/end")
        async def session_end(authorization: Optional[str] = Header(default=None)):
            token = _extract_token(authorization)
            destroyed = await self.session_store.destroy(token) if token else False
            return {"destroyed": destroyed}

        # ============ 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, authorization: Optional[str] = Header(default=None)):
            """
            Stream agent responses via Server-Sent Events (SSE)

            Send ``Authorization: Bearer <token>`` (from /session/start) to keep
            a stable multi-turn session. Without a token an anonymous session is
            used. The backend-generated context_id always wins over the body.

            Expected request format:
            {
                "query": "Your query here",
                "task_id": "optional-task-id"
            }
            """
            query = request.get("query")
            if not query:
                return {"error": "Query is required"}

            context_id = await self._resolve_context_id(authorization)
            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, authorization: Optional[str] = Header(default=None)):
            """
            Single query endpoint (non-streaming)

            Send ``Authorization: Bearer <token>`` (from /session/start) to keep
            a stable multi-turn session; otherwise an anonymous session is used.

            Expected request format:
            {
                "query": "Your query here",
                "task_id": "optional-task-id"
            }
            """
            query = request.get("query")
            if not query:
                return {"error": "Query is required"}

            context_id = await self._resolve_context_id(authorization)
            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