#!/usr/bin/env python3.12

import importlib
import logging
import regex

from aiohttp import ClientSession, ClientTimeout

from fastapi import FastAPI, Security, Request, HTTPException
from fastapi.security import APIKeyHeader, APIKeyQuery


from pydantic import BaseModel


api_key_header = APIKeyHeader(name="X-Authd-With")

    

class AI(FastAPI):
    """AI Endpoints"""   
    def __init__(self, app: FastAPI, my_util, constants, glob_state):  # pylint: disable=super-init-not-called
        self.app = app
        self.util = my_util
        self.constants = constants
        self.glob_state = glob_state
        self.url_clean_regex = regex.compile(r'^\/ai\/')
        self.endpoints = {
            "ai": self.ai_handler,
            #tbd
            }

        for endpoint, handler in self.endpoints.items():
            app.add_api_route(f"/{endpoint}/{{any:path}}", handler, methods=["POST"])
    
    async def ai_handler(self, request: Request):
        """
        /ai/
        AI Request
        """

        if not self.util.check_key(request.url.path, request.headers.get('X-Authd-With')):
            raise HTTPException(status_code=403, detail="Unauthorized")
        
        
        """
        TODO: Implement Claude
        Currently only routes to local LLM
        """

        local_llm_headers = {
            'Authorization': f'Bearer {self.constants.LOCAL_LLM_KEY}'
        }

        forward_path = self.url_clean_regex.sub('', request.url.path)

        print(f"Original path: {request.url.path}; Forward path: {forward_path}")
        print(f"Request data: {await request.json()}")
        try:
                async with ClientSession() as session:
                    async with await session.post(f'{self.constants.LOCAL_LLM_HOST}/{forward_path}',
                                                  json=await request.json(),
                                                  headers=local_llm_headers,
                                                  timeout=ClientTimeout(connect=15, sock_read=30)) as request:
                           await self.glob_state.increment_counter('ai_requests')
                           response = await request.json()
                           print(f"Response received: {response}")
                           return response
        except Exception as e: # pylint: disable=broad-exception-caught
             logging.error("Error: %s", e)
             return {
                  'err': True,
                  'errorText': 'General Failure'
             }