import os
import logging
import json
import time # Import time module
from aiohttp import web
from aiohttp_cors import setup as cors_setup, ResourceOptions
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
logger = logging.getLogger(__name__)

# Import the agent logic
from src.agents import (
    get_simple_litellm_agent,
    get_pydantic_agent,
    PYDANTIC_AI_AVAILABLE
)
from src.tools import get_all_tools
from src.utils import extract_user_id_from_request

async def chat_completions(request):
    """
    Handles OpenAI-compatible chat completion requests using PydanticAI agent.
    Default endpoint - uses PydanticAI Agent with tools enabled.
    """
    if not PYDANTIC_AI_AVAILABLE:
        return web.json_response({
            "error": {"message": "PydanticAI not available. Install with: pip install pydantic-ai"}
        }, status=503)

    try:
        data = await request.json()
        logger.info(f"[DEFAULT/PYDANTIC_AI] Received chat request")

        # Extract relevant fields from the request
        messages = data.get("messages")
        model = data.get("model", "pydantic")
        stream = data.get("stream", False)
        conversation_id = data.get("conversation_id")
        enable_tools = data.get("enable_tools", True)  # Tools enabled by default

        # Extract user ID from request (supports user_id, user_email, or falls back to default)
        user_id = extract_user_id_from_request(data)

        if not messages:
            raise web.HTTPBadRequest(reason="'messages' field is required")

        # Get the agent instance (default: PydanticAI agent with tools)
        # Note: Agent is cached per user_id, so each user gets their own agent instance with their memory
        agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)

        # For non-streaming requests, collect the full response
        if not stream:
            response_content = await agent.chat_completion(
                messages=messages,
                conversation_id=conversation_id
            )
            return web.json_response({
                "id": f"chatcmpl-{os.urandom(12).hex()}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "pydantic",
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": response_content},
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                },
                "tools_enabled": enable_tools,
                "tools_count": len(agent.tools) if enable_tools else 0
            })
        else:
            # Handle streaming response
            response = web.StreamResponse(
                status=200,
                headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'}
            )
            await response.prepare(request)

            async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
                chunk_type = chunk.get("type", "content")

                if chunk_type == "content":
                    json_chunk = {
                        "id": f"chatcmpl-{os.urandom(12).hex()}",
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "model": "pydantic",
                        "choices": [{
                            "index": 0,
                            "delta": {"content": chunk.get("content", "")},
                            "finish_reason": chunk.get("finish_reason")
                        }]
                    }
                    await response.write(f"data: {json.dumps(json_chunk)}\n\n".encode())

                    if chunk.get("finish_reason") == "stop":
                        break
                elif chunk_type == "error":
                    error_chunk = {
                        "error": {"message": chunk.get("content", "Unknown error")}
                    }
                    await response.write(f"data: {json.dumps(error_chunk)}\n\n".encode())
                    break

            await response.write(b"data: [DONE]\n\n")
            await response.write_eof()
            return response

    except web.HTTPBadRequest as e:
        logger.warning(f"Bad request: {e.reason}")
        return web.json_response({"error": {"message": e.reason}}, status=400)
    except Exception as e:
        logger.exception("[DEFAULT/PYDANTIC_AI] Error during chat completion:")
        return web.json_response({"error": {"message": str(e)}}, status=500)

async def chat_simple(request):
    """
    Handles chat requests using SimpleLiteLLMAgent (no tools).
    Endpoint: POST /v1/chat/simple
    """
    try:
        data = await request.json()
        logger.info(f"[SIMPLE] Received chat request")

        messages = data.get("messages")
        model = data.get("model", "simple")
        stream = data.get("stream", False)
        conversation_id = data.get("conversation_id")

        if not messages:
            raise web.HTTPBadRequest(reason="'messages' field is required")

        # Get SimpleLiteLLM agent
        agent = get_simple_litellm_agent()

        # Non-streaming response
        if not stream:
            response_content = await agent.chat_completion(
                messages=messages,
                conversation_id=conversation_id
            )
            return web.json_response({
                "id": f"chatcmpl-{os.urandom(12).hex()}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "simple",
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": response_content},
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                }
            })
        else:
            # Streaming response
            response = web.StreamResponse(
                status=200,
                headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'}
            )
            await response.prepare(request)

            async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
                json_chunk = {
                    "id": f"chatcmpl-{os.urandom(12).hex()}",
                    "object": "chat.completion.chunk",
                    "created": int(time.time()),
                    "model": "simple",
                    "choices": [{
                        "index": 0,
                        "delta": {"content": chunk.get("content", "")},
                        "finish_reason": chunk.get("finish_reason")
                    }]
                }
                await response.write(f"data: {json.dumps(json_chunk)}\n\n".encode())
                if chunk.get("finish_reason") == "stop":
                    break

            await response.write(b"data: [DONE]\n\n")
            await response.write_eof()
            return response

    except web.HTTPBadRequest as e:
        logger.warning(f"Bad request: {e.reason}")
        return web.json_response({"error": {"message": e.reason}}, status=400)
    except Exception as e:
        logger.exception("[SIMPLE] Error during chat completion:")
        return web.json_response({"error": {"message": str(e)}}, status=500)


async def chat_pydantic(request):
    """
    Handles chat requests using PydanticAI Agent with tools.
    Endpoint: POST /v1/chat/pydantic
    """
    if not PYDANTIC_AI_AVAILABLE:
        return web.json_response({
            "error": {"message": "PydanticAI not available. Install with: pip install pydantic-ai"}
        }, status=503)

    try:
        data = await request.json()
        logger.info(f"[PYDANTIC_AI] Received chat request")

        messages = data.get("messages")
        model = data.get("model", "pydantic")
        stream = data.get("stream", False)
        conversation_id = data.get("conversation_id")
        enable_tools = data.get("enable_tools", True)

        # Extract user ID from request
        user_id = extract_user_id_from_request(data)

        if not messages:
            raise web.HTTPBadRequest(reason="'messages' field is required")

        # Get PydanticAI agent with or without tools
        agent = get_pydantic_agent(discover_tools=enable_tools, user_id=user_id)

        # Non-streaming response
        if not stream:
            response_content = await agent.chat_completion(
                messages=messages,
                conversation_id=conversation_id
            )
            return web.json_response({
                "id": f"chatcmpl-{os.urandom(12).hex()}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": "pydantic",
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": response_content},
                    "finish_reason": "stop"
                }],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0
                },
                "tools_enabled": enable_tools,
                "tools_count": len(agent.tools) if enable_tools else 0
            })
        else:
            # Streaming response
            response = web.StreamResponse(
                status=200,
                headers={'Content-Type': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive'}
            )
            await response.prepare(request)

            async for chunk in agent.chat(messages=messages, conversation_id=conversation_id, stream=True):
                chunk_type = chunk.get("type", "content")

                if chunk_type == "content":
                    json_chunk = {
                        "id": f"chatcmpl-{os.urandom(12).hex()}",
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "model": "pydantic",
                        "choices": [{
                            "index": 0,
                            "delta": {"content": chunk.get("content", "")},
                            "finish_reason": chunk.get("finish_reason")
                        }]
                    }
                    await response.write(f"data: {json.dumps(json_chunk)}\n\n".encode())

                    if chunk.get("finish_reason") == "stop":
                        break
                elif chunk_type == "error":
                    error_chunk = {
                        "error": {"message": chunk.get("content", "Unknown error")}
                    }
                    await response.write(f"data: {json.dumps(error_chunk)}\n\n".encode())
                    break

            await response.write(b"data: [DONE]\n\n")
            await response.write_eof()
            return response

    except web.HTTPBadRequest as e:
        logger.warning(f"Bad request: {e.reason}")
        return web.json_response({"error": {"message": e.reason}}, status=400)
    except Exception as e:
        logger.exception("[PYDANTIC_AI] Error during chat completion:")
        return web.json_response({"error": {"message": str(e)}}, status=500)


async def list_models(request):
    """
    Lists available models (OpenAI-compatible endpoint).
    Endpoint: GET /v1/models
    """
    models = [
        {
            "id": "Tatlock",
            "object": "model",
            "created": int(time.time()),
            "owned_by": "core-ai",
            "permission": [],
            "root": "tatlock",
            "parent": None,
        },
        {
            "id": "simple",
            "object": "model",
            "created": int(time.time()),
            "owned_by": "core-ai",
            "permission": [],
            "root": "simple",
            "parent": None,
        }
    ]

    return web.json_response({
        "object": "list",
        "data": models
    })


async def list_tools(request):
    """
    Lists all available tools.
    Endpoint: GET /v1/tools
    """
    try:
        tools = get_all_tools()

        tools_info = []
        for name, func in tools.items():
            tools_info.append({
                "name": name,
                "description": func.__doc__.strip() if func.__doc__ else "No description available",
                "type": "local"
            })

        return web.json_response({
            "tools": tools_info,
            "count": len(tools_info),
            "pydantic_ai_available": PYDANTIC_AI_AVAILABLE
        })

    except Exception as e:
        logger.exception("Error listing tools:")
        return web.json_response({"error": {"message": str(e)}}, status=500)


async def health_check(request):
    """Simple health check endpoint."""
    return web.json_response({
        "status": "ok",
        "service": "core-ai",
        "agents": {
            "simple": True,
            "pydantic": PYDANTIC_AI_AVAILABLE
        },
        "default_agent": "pydantic" if PYDANTIC_AI_AVAILABLE else "simple",
        "tools_count": len(get_all_tools())
    })

async def setup_routes(app):
    # Chat endpoints
    app.router.add_post("/chat/completions", chat_completions)     # Alias without /v1 for compatibility
    app.router.add_post("/v1/chat/completions", chat_completions)  # Default (PydanticAI)
    app.router.add_post("/v1/chat/simple", chat_simple)            # Simple agent (no tools)
    app.router.add_post("/v1/chat/pydantic", chat_pydantic)        # Alias for default

    # OpenAI-compatible endpoints
    app.router.add_get("/v1/models", list_models)                  # List available models
    app.router.add_get("/models", list_models)                     # Alias without /v1 prefix

    # Tool management
    app.router.add_get("/v1/tools", list_tools)                    # List available tools

    # Health check
    app.router.add_get("/health", health_check)

    # Setup CORS
    cors = cors_setup(app, defaults={
        "*": ResourceOptions(
            allow_credentials=True,
            expose_headers="*",
            allow_headers="*",
            allow_methods="*"
        )
    })

    # Configure CORS on all routes
    for route in list(app.router.routes()):
        cors.add(route)

def main():
    app = web.Application()
    app.on_startup.append(setup_routes) # Register routes on startup
    
    # Configuration
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", 8086)) # Use 8086 to avoid conflict with core-ai
    
    logger.info(f"Starting core-ai service on http://{host}:{port}")
    web.run_app(app, host=host, port=port)

if __name__ == "__main__":
    main()