LLM_TIMEOUT = 600

from utils import logger
from pydantic import BaseModel, Field, ValidationError # type: ignore
from pydantic_core import PydanticSerializationError  # type: ignore
from typing import List

from typing import AsyncGenerator, Dict, Optional

# %%
# Imports [standard]
# Standard library modules (no try-except needed)
import argparse
import asyncio
import json
import logging
import os
import re
import uuid
import subprocess
import re
import math
import warnings
from typing import Any
from datetime import datetime
import inspect
from uuid import uuid4
import time
import traceback


def try_import(module_name, pip_name=None):
    try:
        __import__(module_name)
    except ImportError:
        print(f"Module '{module_name}' not found. Install it using:")
        print(f"  pip install {pip_name or module_name}")


# Third-party modules with import checks
try_import("ollama")
try_import("requests")
try_import("fastapi")
try_import("uvicorn")
try_import("numpy")
try_import("umap")
try_import("sklearn")
try_import("prometheus_client")
try_import("prometheus_fastapi_instrumentator")

import ollama
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request  # type: ignore
from fastapi.responses import JSONResponse, StreamingResponse, FileResponse, RedirectResponse  # type: ignore
from fastapi.middleware.cors import CORSMiddleware  # type: ignore
import uvicorn  # type: ignore
import numpy as np  # type: ignore

# Prometheus
from prometheus_client import Summary  # type: ignore
from prometheus_fastapi_instrumentator import Instrumentator  # type: ignore
from prometheus_client import CollectorRegistry, Counter  # type: ignore

from utils import (
    rag as Rag,
    ChromaDBGetResponse,
    tools as Tools,
    Context,
    Conversation,
    Message,
    Agent,
    Metrics,
    Tunables,
    defines,
    check_serializable,
    logger,
)


rags : List[ChromaDBGetResponse] = [
    ChromaDBGetResponse(
        name="JPK",
        enabled=True,
        description="Expert data about James Ketrenos, including work history, personal hobbies, and projects.",
    ),
    # { "name": "LKML", "enabled": False, "description": "Full associative data for entire LKML mailing list archive." },
]


class Query(BaseModel):
    prompt: str
    tunables: Tunables = Field(default_factory=Tunables)
    agent_options: Dict[str, Any] = Field(default={})


REQUEST_TIME = Summary("request_processing_seconds", "Time spent processing request")


def get_installed_ram():
    try:
        with open("/proc/meminfo", "r") as f:
            meminfo = f.read()
        match = re.search(r"MemTotal:\s+(\d+)", meminfo)
        if match:
            return f"{math.floor(int(match.group(1)) / 1000**2)}GB"  # Convert KB to GB
    except Exception as e:
        return f"Error retrieving RAM: {e}"


def get_graphics_cards():
    gpus = []
    try:
        # Run the ze-monitor utility
        result = subprocess.run(
            ["ze-monitor"], capture_output=True, text=True, check=True
        )

        # Clean up the output (remove leading/trailing whitespace and newlines)
        output = result.stdout.strip()
        for index in range(len(output.splitlines())):
            result = subprocess.run(
                ["ze-monitor", "--device", f"{index+1}", "--info"],
                capture_output=True,
                text=True,
                check=True,
            )
            gpu_info = result.stdout.strip().splitlines()
            gpu = {
                "discrete": True,  # Assume it's discrete initially
                "name": None,
                "memory": None,
            }
            gpus.append(gpu)
            for line in gpu_info:
                match = re.match(r"^Device: [^(]*\((.*)\)", line)
                if match:
                    gpu["name"] = match.group(1)
                    continue

                match = re.match(r"^\s*Memory: (.*)", line)
                if match:
                    gpu["memory"] = match.group(1)
                    continue

                match = re.match(r"^.*Is integrated with host: Yes.*", line)
                if match:
                    gpu["discrete"] = False
                    continue

        return gpus
    except Exception as e:
        return f"Error retrieving GPU info: {e}"


def get_cpu_info():
    try:
        with open("/proc/cpuinfo", "r") as f:
            cpuinfo = f.read()
        model_match = re.search(r"model name\s+:\s+(.+)", cpuinfo)
        cores_match = re.findall(r"processor\s+:\s+\d+", cpuinfo)
        if model_match and cores_match:
            return f"{model_match.group(1)} with {len(cores_match)} cores"
    except Exception as e:
        return f"Error retrieving CPU info: {e}"

def system_info(model):
    return {
        "System RAM": get_installed_ram(),
        "Graphics Card": get_graphics_cards(),
        "CPU": get_cpu_info(),
        "LLM Model": model,
        "Embedding Model": defines.embedding_model,
        "Context length": defines.max_context,
    }

# %%
# Defaults
OLLAMA_API_URL = defines.ollama_api_url
MODEL_NAME = defines.model
LOG_LEVEL = "info"
WEB_HOST = "0.0.0.0"
WEB_PORT = 8911
DEFAULT_HISTORY_LENGTH = 5

# %%
# Globals

model = None
web_server = None


# %%
# Cmd line overrides
def parse_args():
    parser = argparse.ArgumentParser(description="AI is Really Cool")
    parser.add_argument(
        "--ollama-server",
        type=str,
        default=OLLAMA_API_URL,
        help=f"Ollama API endpoint. default={OLLAMA_API_URL}",
    )
    parser.add_argument(
        "--ollama-model",
        type=str,
        default=MODEL_NAME,
        help=f"LLM model to use. default={MODEL_NAME}",
    )
    parser.add_argument(
        "--web-host",
        type=str,
        default=WEB_HOST,
        help=f"Host to launch Flask web server. default={WEB_HOST} only if --web-disable not specified.",
    )
    parser.add_argument(
        "--web-port",
        type=str,
        default=WEB_PORT,
        help=f"Port to launch Flask web server. default={WEB_PORT} only if --web-disable not specified.",
    )
    parser.add_argument(
        "--level",
        type=str,
        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        default=LOG_LEVEL,
        help=f"Set the logging level. default={LOG_LEVEL}",
    )
    return parser.parse_args()


# %%


# %%


# %%
def is_valid_uuid(value: str) -> bool:
    try:
        uuid_obj = uuid.UUID(value, version=4)
        return str(uuid_obj) == value
    except (ValueError, TypeError):
        return False


# %%
class WebServer:
    @asynccontextmanager
    async def lifespan(self, app: FastAPI):
        # Start the file watcher
        self.observer, self.file_watcher = Rag.start_file_watcher(
            llm=self.llm,
            watch_directory=defines.doc_dir,
            recreate=False,  # Don't recreate if exists
        )
        logger.info(
            f"API started with {self.file_watcher.collection.count()} documents in the collection"
        )
        yield
        if self.observer:
            self.observer.stop()
            self.observer.join()
        logger.info("File watcher stopped")

    def __init__(self, llm, model=MODEL_NAME):
        self.app = FastAPI(lifespan=self.lifespan)

        self.prometheus_collector = CollectorRegistry()
        self.metrics = Metrics(prometheus_collector=self.prometheus_collector)

        # Keep the Instrumentator instance alive
        self.instrumentator = Instrumentator(registry=self.prometheus_collector)

        # Instrument the FastAPI app
        self.instrumentator.instrument(self.app)

        # Expose the /metrics endpoint
        self.instrumentator.expose(self.app, endpoint="/metrics")

        self.contexts = {}

        self.llm = llm
        self.model = model
        self.processing = False
        self.file_watcher = None
        self.observer = None

        self.ssl_enabled = os.path.exists(defines.key_path) and os.path.exists(
            defines.cert_path
        )

        if self.ssl_enabled:
            allow_origins = ["https://battle-linux.ketrenos.com:3000"]
        else:
            allow_origins = ["http://battle-linux.ketrenos.com:3000"]

        logger.info(f"Allowed origins: {allow_origins}")

        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=allow_origins,
            allow_credentials=True,
            allow_methods=["*"],
            allow_headers=["*"],
        )

        self.setup_routes()

    def setup_routes(self):
        @self.app.get("/")
        async def root():
            context = self.create_context()
            logger.info(f"Redirecting non-context to {context.id}")
            return RedirectResponse(url=f"/{context.id}", status_code=307)
            # return JSONResponse({"redirect": f"/{context.id}"})

        @self.app.get("/api/umap/entry/{doc_id}/{context_id}")
        async def get_umap(doc_id: str, context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            try:
                if not self.file_watcher:
                    raise Exception("File watcher not initialized")

                context = self.upsert_context(context_id)
                if not context:
                    return JSONResponse(
                        {"error": f"Invalid context: {context_id}"}, status_code=400
                    )
                collection = self.file_watcher.umap_collection
                if not collection:
                    return JSONResponse(
                        {"error": "No UMAP collection found"}, status_code=404
                    )
                if not collection.get("metadatas", None):
                    return JSONResponse(f"Document id {doc_id} not found.", 404)

                for index, id in enumerate(collection.get("ids", [])):
                    if id == doc_id:
                        metadata = collection.get("metadatas", [])[index].copy()
                        content = self.file_watcher.prepare_metadata(metadata)
                        return JSONResponse(content)
                    
                return JSONResponse(f"Document id {doc_id} not found.", 404)
            except Exception as e:
                logger.error(f"get_umap error: {str(e)}")
                logger.error(traceback.format_exc())
                return JSONResponse({"error": str(e)}, 500)

        @self.app.put("/api/umap/{context_id}")
        async def put_umap(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            try:
                if not self.file_watcher:
                    raise Exception("File watcher not initialized")

                context = self.upsert_context(context_id)
                if not context:
                    return JSONResponse(
                        {"error": f"Invalid context: {context_id}"}, status_code=400
                    )

                data = await request.json()

                dimensions = data.get("dimensions", 2)
                collection = self.file_watcher.umap_collection
                if not collection:
                    return JSONResponse(
                        {"error": "No UMAP collection found"}, status_code=404
                    )
                if dimensions == 2:
                    logger.info("Returning 2D UMAP")
                    umap_embedding = self.file_watcher.umap_embedding_2d
                else:
                    logger.info("Returning 3D UMAP")
                    umap_embedding = self.file_watcher.umap_embedding_3d

                if len(umap_embedding) == 0:
                    return JSONResponse(
                        {"error": "No UMAP embedding found"}, status_code=404
                    )
                result = {
                    "ids": collection.get("ids", []),
                    "metadatas": collection.get("metadatas", []),
                    "documents": collection.get("documents", []),
                    "embeddings": umap_embedding.tolist(),
                    "size": self.file_watcher.collection.count()
                }

                return JSONResponse(result)

            except Exception as e:
                logger.error(f"put_umap error: {str(e)}")
                logger.error(traceback.format_exc())
                return JSONResponse({"error": str(e)}, 500)

        @self.app.put("/api/similarity/{context_id}")
        async def put_similarity(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            if not self.file_watcher:
                raise Exception("File watcher not initialized")

            if not is_valid_uuid(context_id):
                logger.warning(f"Invalid context_id: {context_id}")
                return JSONResponse({"error": "Invalid context_id"}, status_code=400)

            try:
                data = await request.json()
                query = data.get("query", "")
                threshold = data.get("threshold", defines.default_rag_threshold)
                results = data.get("results", defines.default_rag_top_k)
            except:
                query = ""
                threshold = defines.default_rag_threshold
                results = defines.default_rag_top_k
            if not query:
                return JSONResponse(
                    {"error": "No query provided for similarity search"},
                    status_code=400,
                )
            try:
                chroma_results = self.file_watcher.find_similar(
                    query=query, top_k=results, threshold=threshold
                )
                if not chroma_results:
                    return JSONResponse({"error": "No results found"}, status_code=404)

                chroma_embedding = np.array(
                    chroma_results["query_embedding"]
                ).flatten()  # Ensure correct shape
                logger.info(f"Chroma embedding shape: {chroma_embedding.shape}")

                umap_2d = self.file_watcher.umap_model_2d.transform([chroma_embedding])[
                    0
                ].tolist()
                logger.info(
                    f"UMAP 2D output: {umap_2d}, length: {len(umap_2d)}"
                )  # Debug output

                umap_3d = self.file_watcher.umap_model_3d.transform([chroma_embedding])[
                    0
                ].tolist()
                logger.info(
                    f"UMAP 3D output: {umap_3d}, length: {len(umap_3d)}"
                )  # Debug output

                return JSONResponse({
                    "distances": chroma_results["distances"],
                    "ids": chroma_results["ids"],
                    "metadatas": chroma_results["metadatas"],
                    "query": query,
                    "umap_embedding_2d": umap_2d,
                    "umap_embedding_3d": umap_3d,
                    "size": self.file_watcher.collection.count()
                })

            except Exception as e:
                logger.error(e)
                logging.error(traceback.format_exc())
                return JSONResponse({"error": str(e)}, 500)

        @self.app.put("/api/reset/{context_id}/{agent_type}")
        async def put_reset(context_id: str, agent_type: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            if not is_valid_uuid(context_id):
                logger.warning(f"Invalid context_id: {context_id}")
                return JSONResponse({"error": "Invalid context_id"}, status_code=400)
            context = self.upsert_context(context_id)
            agent = context.get_agent(agent_type)
            if not agent:
                response = { "history": [] }
                return JSONResponse(response)

            data = await request.json()
            try:
                response = {}
                for reset_operation in data["reset"]:
                    match reset_operation:
                        case "system_prompt":
                            logger.info(f"Resetting {reset_operation}")
                        case "rags":
                            logger.info(f"Resetting {reset_operation}")
                            context.rags = [ r.model_copy() for r in rags]
                            response["rags"] = [ r.model_dump(mode="json") for r in context.rags ]
                        case "tools":
                            logger.info(f"Resetting {reset_operation}")
                            context.tools = Tools.enabled_tools(Tools.tools)
                            response["tools"] = context.tools
                        case "history":
                            reset_map = {
                                "job_description": (
                                    "job_description",
                                    "resume",
                                    "fact_check",
                                ),
                                "resume": ("job_description", "resume", "fact_check"),
                                "fact_check": (
                                    "job_description",
                                    "resume",
                                    "fact_check",
                                ),
                                "chat": ("chat",),
                            }
                            resets = reset_map.get(agent_type, ())
                            for mode in resets:
                                tmp = context.get_agent(mode)
                                if not tmp:
                                    logger.info(
                                        f"Agent {mode} not found for {context_id}"
                                    )
                                    continue
                                logger.info(f"Resetting {reset_operation} for {mode}")
                                tmp.conversation.reset()
                            response["history"] = []
                            response["context_used"] = agent.context_tokens
                        case "message_history_length":
                            logger.info(f"Resetting {reset_operation}")
                            context.message_history_length = DEFAULT_HISTORY_LENGTH
                            response["message_history_length"] = DEFAULT_HISTORY_LENGTH

                if not response:
                    return JSONResponse(
                        {"error": "Usage: { reset: rags|tools|history|system_prompt}"}
                    )
                else:
                    self.save_context(context_id)
                    return JSONResponse(response)

            except Exception as e:
                logger.error(f"Error in reset: {e}")
                logger.error(traceback.format_exc())
                return JSONResponse(
                    {"error": "Usage: { reset: rags|tools|history|system_prompt}"}
                )

        @self.app.put("/api/tunables/{context_id}")
        async def put_tunables(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            try:
                context = self.upsert_context(context_id)

                data = await request.json()
                agent = context.get_agent("chat")
                if not agent:
                    logger.info("chat agent does not exist on this context!")
                    return JSONResponse(
                        {"error": f"chat is not recognized", "context": context.id},
                        status_code=404,
                    )
                for k in data.keys():
                    match k:
                        case "tools":
                            # { "tools": [{ "tool": tool?.name, "enabled": tool.enabled }] }
                            tools: list[dict[str, Any]] = data[k]
                            if not tools:
                                return JSONResponse(
                                    {
                                        "status": "error",
                                        "message": "Tools can not be empty.",
                                    }
                                )
                            for tool in tools:
                                for context_tool in context.tools:
                                    if context_tool["function"]["name"] == tool["name"]:
                                        context_tool["enabled"] = tool["enabled"]
                            self.save_context(context_id)
                            return JSONResponse(
                                {
                                    "tools": [
                                        {
                                            **t["function"],
                                            "enabled": t["enabled"],
                                        }
                                        for t in context.tools
                                    ]
                                }
                            )

                        case "rags":
                            # { "rags": [{ "tool": tool?.name, "enabled": tool.enabled }] }
                            rag_configs: list[dict[str, Any]] = data[k]
                            if not rag_configs:
                                return JSONResponse(
                                    {
                                        "status": "error",
                                        "message": "RAGs can not be empty.",
                                    }
                                )
                            for config in rag_configs:
                                for context_rag in context.rags:
                                    if context_rag.name == config["name"]:
                                        context_rag.enabled = config["enabled"]
                            self.save_context(context_id)
                            return JSONResponse({"rags": [ r.model_dump(mode="json") for r in context.rags]})

                        case "system_prompt":
                            system_prompt = data[k].strip()
                            if not system_prompt:
                                return JSONResponse(
                                    {
                                        "status": "error",
                                        "message": "System prompt can not be empty.",
                                    }
                                )
                            agent.system_prompt = system_prompt
                            self.save_context(context_id)
                            return JSONResponse({"system_prompt": system_prompt})
                        case "message_history_length":
                            value = max(0, int(data[k]))
                            context.message_history_length = value
                            self.save_context(context_id)
                            return JSONResponse({"message_history_length": value})
                        case _:
                            return JSONResponse(
                                {"error": f"Unrecognized tunable {k}"}, status_code=404
                            )
            except Exception as e:
                logger.error(f"Error in put_tunables: {e}")
                return JSONResponse({"error": str(e)}, status_code=500)

        @self.app.get("/api/tunables/{context_id}")
        async def get_tunables(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            context = self.upsert_context(context_id)
            agent = context.get_agent("chat")
            if not agent:
                logger.info("chat agent does not exist on this context!")
                return JSONResponse(
                    {"error": f"chat is not recognized", "context": context.id},
                    status_code=404,
                )
            return JSONResponse(
                {
                    "system_prompt": agent.system_prompt,
                    "message_history_length": context.message_history_length,
                    "rags": [ r.model_dump(mode="json") for r in context.rags ],
                    "tools": [
                        {
                            **t["function"],
                            "enabled": t["enabled"],
                        }
                        for t in context.tools
                    ],
                }
            )

        @self.app.get("/api/system-info/{context_id}")
        async def get_system_info(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            return JSONResponse(system_info(self.model))

        @self.app.post("/api/{agent_type}/{context_id}")
        async def post_agent_endpoint(
            agent_type: str, context_id: str, request: Request
        ):
            logger.info(f"{request.method} {request.url.path}")

            try:
                context = self.upsert_context(context_id)
            except Exception as e:
                error = {
                    "error": f"Unable to create or access context {context_id}: {e}"
                }
                logger.info(error)
                return JSONResponse(error, status_code=404)

            try:
                data = await request.json()
                query: Query = Query(**data)
            except Exception as e:
                error = {"error": f"Attempt to parse request: {e}"}
                logger.info(error)
                return JSONResponse(error, status_code=400)

            try:
                agent = context.get_or_create_agent(agent_type, **query.agent_options)
            except Exception as e:
                error = {
                    "error": f"Attempt to create agent type: {agent_type} failed: {e}"
                }
                logger.info(error)
                return JSONResponse(error, status_code=404)

            try:

                async def flush_generator():
                    logger.info(f"{agent.agent_type} - {inspect.stack()[0].function}")
                    try:
                        start_time = time.perf_counter()
                        async for message in self.generate_response(
                            context=context,
                            agent=agent,
                            prompt=query.prompt,
                            tunables=query.tunables,
                        ):
                            if message.status != "done" and message.status != "partial":
                                if message.status == "streaming":
                                    result = {
                                        "status": "streaming",
                                        "chunk": message.chunk,
                                        "remaining_time": LLM_TIMEOUT
                                        - (time.perf_counter() - start_time),
                                    }
                                else:
                                    start_time = time.perf_counter()
                                    result = {
                                        "status": message.status,
                                        "response": message.response,
                                        "remaining_time": LLM_TIMEOUT,
                                    }
                            else:
                                logger.info(f"Providing {message.status} response.")
                                try:
                                    result = message.model_dump(
                                        by_alias=True, mode="json"
                                    )
                                except Exception as e:
                                    result = {"status": "error", "response": str(e)}
                                    yield json.dumps(result) + "\n"
                                    return

                            # Convert to JSON and add newline
                            result = json.dumps(result) + "\n"
                            message.network_packets += 1
                            message.network_bytes += len(result)
                            yield result
                            if await request.is_disconnected():
                                logger.info("Disconnect detected. Aborting generation.")
                                context.processing = False
                                # Save context on completion or error
                                message.prompt = query.prompt
                                message.status = "error"
                                message.response = (
                                    "Client disconnected during generation."
                                )
                                agent.conversation.add(message)
                                self.save_context(context_id)
                                return

                            current_time = time.perf_counter()
                            if current_time - start_time > LLM_TIMEOUT:
                                message.status = "error"
                                message.response = f"Processing time ({LLM_TIMEOUT}s) exceeded for single LLM inference (likely due to LLM getting stuck.) You will need to retry your query."
                                message.partial_response = message.response
                                logger.info(message.response + " Ending session")
                                result = message.model_dump(by_alias=True, mode="json")
                                result = json.dumps(result) + "\n"
                                yield result

                            if message.status == "error":
                                context.processing = False
                                return

                            # Allow the event loop to process the write
                            await asyncio.sleep(0)
                    except Exception as e:
                        context.processing = False
                        logger.error(f"Error in generate_response: {e}")
                        logger.error(traceback.format_exc())
                        yield json.dumps({"status": "error", "response": str(e)}) + "\n"
                    finally:
                        # Save context on completion or error
                        self.save_context(context_id)

                # Return StreamingResponse with appropriate headers
                return StreamingResponse(
                    flush_generator(),
                    media_type="application/json",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-Accel-Buffering": "no",  # Prevents Nginx buffering if you're using it
                    },
                )
            except Exception as e:
                context.processing = False
                logger.error(f"Error in post_chat_endpoint: {e}")
                return JSONResponse({"error": str(e)}, status_code=500)

        @self.app.post("/api/context")
        async def create_context():
            try:
                context = self.create_context()
                logger.info(f"Generated new agent as {context.id}")
                return JSONResponse({"id": context.id})
            except Exception as e:
                logger.error(f"get_history error: {str(e)}")
                logger.error(traceback.format_exc())
                return JSONResponse({"error": str(e)}, status_code=404)

        @self.app.get("/api/history/{context_id}/{agent_type}")
        async def get_history(context_id: str, agent_type: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            try:
                context = self.upsert_context(context_id)
                agent = context.get_agent(agent_type)
                if not agent:
                    logger.info(
                        f"Agent {agent_type} not found. Returning empty history."
                    )
                    return JSONResponse({"messages": []})
                logger.info(
                    f"History for {agent_type} contains {len(agent.conversation)} entries."
                )
                return agent.conversation
            except Exception as e:
                logger.error(f"get_history error: {str(e)}")
                import traceback

                logger.error(traceback.format_exc())
                return JSONResponse({"error": str(e)}, status_code=404)

        @self.app.get("/api/tools/{context_id}")
        async def get_tools(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            context = self.upsert_context(context_id)
            return JSONResponse(context.tools)

        @self.app.put("/api/tools/{context_id}")
        async def put_tools(context_id: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            if not is_valid_uuid(context_id):
                logger.warning(f"Invalid context_id: {context_id}")
                return JSONResponse({"error": "Invalid context_id"}, status_code=400)
            context = self.upsert_context(context_id)
            try:
                data = await request.json()
                modify = data["tool"]
                enabled = data["enabled"]
                for tool in context.tools:
                    if modify == tool["function"]["name"]:
                        tool["enabled"] = enabled
                        self.save_context(context_id)
                        return JSONResponse(context.tools)
                return JSONResponse(
                    {"status": f"{modify} not found in tools."}, status_code=404
                )
            except:
                return JSONResponse({"status": "error"}, 405)

        @self.app.get("/api/context-status/{context_id}/{agent_type}")
        async def get_context_status(context_id, agent_type: str, request: Request):
            logger.info(f"{request.method} {request.url.path}")
            if not is_valid_uuid(context_id):
                logger.warning(f"Invalid context_id: {context_id}")
                return JSONResponse({"error": "Invalid context_id"}, status_code=400)
            context = self.upsert_context(context_id)
            agent = context.get_agent(agent_type)
            if not agent:
                return JSONResponse(
                    {"context_used": 0, "max_context": defines.max_context}
                )
            return JSONResponse(
                {
                    "context_used": agent.context_tokens,
                    "max_context": defines.max_context,
                }
            )

        @self.app.get("/api/health")
        async def health_check():
            return JSONResponse({"status": "healthy"})

        @self.app.get("/{path:path}")
        async def serve_static(path: str, request: Request):
            full_path = os.path.join(defines.static_content, path)
            if os.path.exists(full_path) and os.path.isfile(full_path):
                logger.info(f"Serve static request for {full_path}")
                return FileResponse(full_path)

            logger.info(f"Serve index.html for {path}")
            return FileResponse(os.path.join(defines.static_content, "index.html"))

    def save_context(self, context_id):
        """
        Serialize a Python dictionary to a file in the agents directory.

        Args:
            data: Dictionary containing the agent data
            context_id: UUID string for the context. If it doesn't exist, it is created

        Returns:
            The context_id used for the file
        """
        context = self.upsert_context(context_id)

        # Create agents directory if it doesn't exist
        if not os.path.exists(defines.context_dir):
            os.makedirs(defines.context_dir)

        # Create the full file path
        file_path = os.path.join(defines.context_dir, context_id)

        # Serialize the data to JSON and write to file
        try:
            # Check for non-serializable fields before dumping
            serialization_errors = check_serializable(context)
            if serialization_errors:
                for error in serialization_errors:
                    logger.error(error)
                raise ValueError("Found non-serializable fields in the model")
            # Dump the model prior to opening file in case there is
            # a validation error so it doesn't delete the current
            # context session
            json_data = context.model_dump_json(by_alias=True)
            with open(file_path, "w") as f:
                f.write(json_data)
        except ValidationError as e:
            logger.error(e)
            logger.error(traceback.format_exc())
            for error in e.errors():
                print(f"Field: {error['loc'][0]}, Error: {error['msg']}")
        except PydanticSerializationError as e:
            logger.error(e)
            logger.error(traceback.format_exc())
            logger.error(f"Serialization error: {str(e)}")
            # Inspect the model to identify problematic fields
            for field_name, value in context.__dict__.items():
                if isinstance(value, np.ndarray):
                    logger.error(f"Field '{field_name}' contains non-serializable type: {type(value)}")                            
        except Exception as e:
            logger.error(traceback.format_exc())
            logger.error(e)

        return context_id

    def load_or_create_context(self, context_id) -> Context:
        """
        Load a context from a file in the context directory or create a new one if it doesn't exist.
        Args:
            context_id: UUID string for the context.
        Returns:
            A Context object with the specified ID and default settings.
        """
        if not self.file_watcher:
            raise Exception("File watcher not initialized")

        file_path = os.path.join(defines.context_dir, context_id)

        # Check if the file exists
        if not os.path.exists(file_path):
            logger.info(f"Context file {file_path} not found. Creating new context.")
            self.contexts[context_id] = self.create_context(context_id)
        else:
            # Read and deserialize the data
            with open(file_path, "r") as f:
                content = f.read()
                logger.info(
                    f"Loading context from {file_path}, content length: {len(content)}"
                )
                import json

                try:
                    # Try parsing as JSON first to ensure valid JSON
                    json_data = json.loads(content)
                    logger.info("JSON parsed successfully, attempting model validation")

                    # Validate from JSON (no prometheus_collector or file_watcher)
                    context = Context.model_validate(json_data)

                    # Set excluded fields
                    context.file_watcher = self.file_watcher
                    context.prometheus_collector = self.prometheus_collector

                    # Now set context on agents manually
                    agent_types = [agent.agent_type for agent in context.agents]
                    if len(agent_types) != len(set(agent_types)):
                        raise ValueError(
                            "Context cannot contain multiple agents of the same agent_type"
                        )
                    for agent in context.agents:
                        agent.set_context(context)

                    self.contexts[context_id] = context

                    logger.info(f"Successfully loaded context {context_id}")
                except Exception as e:
                    logger.error(f"Error validating context: {str(e)}")
                    logger.error(traceback.format_exc())
                    # Fallback to creating a new context
                    self.contexts[context_id] = Context(
                        id=context_id,
                        file_watcher=self.file_watcher,
                        prometheus_collector=self.prometheus_collector,
                    )

        return self.contexts[context_id]

    def create_context(self, context_id=None) -> Context:
        """
        Create a new context with a unique ID and default settings.
        Args:
            context_id: Optional UUID string for the context. If not provided, a new UUID is generated.
        Returns:
            A Context object with the specified ID and default settings.
        """
        if not self.file_watcher:
            raise Exception("File watcher not initialized")
        if not context_id:
            context_id = str(uuid4())
        logger.info(f"Creating new context with ID: {context_id}")
        context = Context(
            id=context_id,
            file_watcher=self.file_watcher,
            prometheus_collector=self.prometheus_collector,
        )

        if os.path.exists(defines.resume_doc):
            context.user_resume = open(defines.resume_doc, "r").read()
        context.get_or_create_agent(agent_type="chat")
        #    system_prompt=system_message)
        # context.add_agent(Resume(system_prompt = system_generate_resume))
        # context.add_agent(JobDescription(system_prompt = system_job_description))
        # context.add_agent(FactCheck(system_prompt = system_fact_check))
        context.tools = Tools.enabled_tools(Tools.tools)
        context.rags = [ r.model_copy() for r in rags ]

        logger.info(f"{context.id} created and added to contexts.")
        self.contexts[context.id] = context
        self.save_context(context.id)
        return context

    def upsert_context(self, context_id=None) -> Context:
        """
        Upsert a context based on the provided context_id.
        Args:
            context_id: UUID string for the context. If it doesn't exist, a new context is created.
        Returns:
            A Context object with the specified ID and default settings.
        """

        if not context_id:
            logger.warning("No context ID provided. Creating a new context.")
            return self.create_context()

        if context_id in self.contexts:
            return self.contexts[context_id]

        logger.info(f"Context {context_id} is not yet loaded.")
        return self.load_or_create_context(context_id)

    @REQUEST_TIME.time()
    async def generate_response(
        self, context: Context, agent: Agent, prompt: str, tunables: Tunables | None
    ) -> AsyncGenerator[Message, None]:
        if not self.file_watcher:
            raise Exception("File watcher not initialized")

        agent_type = agent.get_agent_type()
        logger.info(f"generate_response: type - {agent_type}")
        # Merge tunables to take agent defaults and override with user supplied settings
        agent_tunables = agent.tunables.model_dump() if agent.tunables else {}
        user_tunables = tunables.model_dump() if tunables else {}
        merged_tunables = {**agent_tunables, **user_tunables}
        message = Message(prompt=prompt, tunables=Tunables(**merged_tunables))

        async for message in agent.prepare_message(message):
            # logger.info(f"{agent_type}.prepare_message: {value.status} - {value.response}")
            if message.status == "error":
                yield message
                return
            if message.status != "done":
                yield message

        async for message in agent.process_message(self.llm, self.model, message):
            if message.status != "done":
                yield message

        if message.status == "error":
            return

        logger.info(
            f"{agent_type}.process_message: {message.status} {f'...{message.response[-20:]}' if len(message.response) > 20 else message.response}"
        )
        message.status = "done"
        yield message
        return

    def run(self, host="0.0.0.0", port=WEB_PORT, **kwargs):
        try:
            if self.ssl_enabled:
                logger.info(f"Starting web server at https://{host}:{port}")
                uvicorn.run(
                    self.app,
                    host=host,
                    port=port,
                    log_config=None,
                    ssl_keyfile=defines.key_path,
                    ssl_certfile=defines.cert_path,
                )
            else:
                logger.info(f"Starting web server at http://{host}:{port}")
                uvicorn.run(self.app, host=host, port=port, log_config=None)
        except KeyboardInterrupt:
            if self.observer:
                self.observer.stop()
        if self.observer:
            self.observer.join()


# %%


# Main function to run everything
def main():
    global model

    # Parse command-line arguments
    args = parse_args()

    # Setup logging based on the provided level
    logger.setLevel(args.level.upper())

    warnings.filterwarnings("ignore", category=FutureWarning, module="sklearn.*")

    warnings.filterwarnings("ignore", category=UserWarning, module="umap.*")

    llm = ollama.Client(host=args.ollama_server)  # type: ignore
    model = args.ollama_model

    web_server = WebServer(llm, model)

    web_server.run(host=args.web_host, port=args.web_port, use_reloader=False)


main()