import os
import sys
import glob
import time
import hashlib
import textwrap
import requests
import httpx
import socket
import json
import random
import traceback
from fastapi.responses import StreamingResponse
from fastapi import FastAPI, Query, HTTPException, Request, Body
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from ddgs import DDGS
from unstructured.cleaners.core import clean_extra_whitespace, clean_non_ascii_chars, replace_unicode_quotes
from datetime import datetime, timezone, timedelta
from pydantic import BaseModel, Field, PrivateAttr, ValidationError
from typing import Optional, List, Dict, Any, Union, Literal
from collections import Counter


# --- Configuration via variables d'environnement ---

PERSIST_DIR = os.environ.get("CHROMA_PERSIST_DIR", "/chroma_db")
CACHE_DIR = os.environ.get("RESPONSE_CACHE_DIR", "/response_cache")
GENERATE_MODEL = os.environ.get("GENERATE_MODEL", "deepseek-coder:6.7b-instruct")
CHAT_MODEL = os.environ.get("CHAT_MODEL", GENERATE_MODEL)
EMBED_MODEL = os.environ.get("EMBED_MODEL", "nomic-embed-text:latest")
SRC_PATH=os.environ.get("SRC_PATH", ".")
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://host.docker.internal:11434")
PROG_LANG = os.environ.get("PROG_LANG", "go")

try:
    CHUNCK_SIZE = int(os.environ.get("CHUNCK_SIZE", "300"))
except (ValueError, TypeError):
    CHUNCK_SIZE = 300

try:
    CHUNCK_OVERLAP = int(os.environ.get("CHUNCK_OVERLAP", "50"))
except (ValueError, TypeError):
    CHUNCK_OVERLAP = 50

try:
    QUERY_TIMEOUT = int(os.environ.get("QUERY_TIMEOUT", "120"))
except (ValueError, TypeError):
    QUERY_TIMEOUT = 120

DDGS_SEARCH_ENABLED = True


os.makedirs(PERSIST_DIR, exist_ok=True)
os.makedirs(CACHE_DIR, exist_ok=True)


# --- Nettoyage du code ---
def clean_code_content(content: str) -> str:
    cleaned = replace_unicode_quotes(content)
    cleaned = clean_non_ascii_chars(cleaned)
    cleaned = clean_extra_whitespace(cleaned)
    return cleaned

# --- Cache simple ---
def get_cache_key(question: str) -> str:
    return hashlib.md5(question.encode()).hexdigest()

# --- Hot-reload : hash du code ---
def hash_code_dir(paths: list) -> str:
    m = hashlib.md5()
    for path in paths:
        abs_path = os.path.join("/code", path) if path != "." else "/code"
        for f in glob.glob(f"{abs_path}/**/*.go", recursive=True):
            try:
                with open(f, "rb") as file:
                    m.update(file.read())
            except Exception:
                continue
    return m.hexdigest()

# --- Wrapper Nomic Embeddings ---
from typing import List

class NomicEmbeddingsWrapper(OllamaEmbeddings):
    """Wrapper Ollama pour les embeddings de code, compatible Chroma."""
    _cached_dim: int = PrivateAttr()  # attribut interne non validé par Pydantic

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # Calcul de la dimension une seule fois
        self._cached_dim = 768 #len(super().embed_query("Hello"))

    def _prefix_text(self, text: str, is_document: bool) -> str:
        """Ajoute un préfixe pour distinguer document vs query."""
        prefix = "search_document: " if is_document else "search_query: "
        return prefix + text

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embeds documents en ajoutant le préfixe, retourne liste de vecteurs float."""
        prefixed_texts = [self._prefix_text(t, is_document=True) for t in texts]
        embeddings = super().embed_documents(prefixed_texts)
        # Normaliser les embeddings vides pour éviter les erreurs Chroma
        return [e if e else [0.0] * self.model_dimensions for e in embeddings]

    def embed_query(self, text: str) -> List[float]:
        """Embeds une query en ajoutant le préfixe."""
        emb = super().embed_query(self._prefix_text(text, is_document=False))
        # Normaliser embedding vide
        return emb if emb else [0.0] * self.model_dimensions

    @property
    def model_dimensions(self) -> int:
        return self._cached_dim

# --- FastAPI ---
app = FastAPI()

# --- Traitement des chemins ---


src_paths_directories = SRC_PATH.split(":")
if not src_paths_directories:
    src_paths_directories = ["."]

# --- Initialisation ---
vectorstore = None
code_hash = ""

def build_vectorstore():
    global vectorstore, code_hash, src_paths_directories
    print("🔹 Construction du vectorstore ...", file=sys.stderr)
    

    # Hash du code pour hot-reload
    new_hash = hash_code_dir(src_paths_directories)
    if vectorstore and new_hash == code_hash:
        print("🔹 Pas de changement dans /code, utilisation du vectorstore existant", file=sys.stderr)
        return
    code_hash = new_hash

    print(f"  🔹 Programmation language: {PROG_LANG}", file=sys.stderr)
    print(f"  🔹 Chunck size: {CHUNCK_SIZE}", file=sys.stderr)
    print(f"  🔹 Chunck overlap: {CHUNCK_OVERLAP}", file=sys.stderr)
    # Text splitter optimisé Go
    code_splitter = RecursiveCharacterTextSplitter.from_language(
        language=PROG_LANG,
        chunk_size=CHUNCK_SIZE,
        chunk_overlap=CHUNCK_OVERLAP, 
        keep_separator=True 
    )

    all_docs = []
    for path in src_paths_directories:
        abs_path = os.path.join("/code", path) if path != "." else "/code"
        print(f"   🔹 Chargement du code Go depuis: {abs_path}", file=sys.stderr)
        loader = DirectoryLoader(
            abs_path,
            glob="**/*.go",
            loader_cls=TextLoader,
            use_multithreading=True,
            loader_kwargs={'autodetect_encoding': True}
        )
        loaded_docs = loader.load()
        print(f"   🔸 {len(loaded_docs)} fichiers chargés", file=sys.stderr)
        for doc in loaded_docs:
            doc.page_content = clean_code_content(doc.page_content)
        all_docs.extend(loaded_docs)

    print(f"🔹 {len(all_docs)} documents après chargement", file=sys.stderr)

    embedding = NomicEmbeddingsWrapper(model=EMBED_MODEL, base_url=OLLAMA_HOST)

    chat_collection = Chroma.from_documents(
        documents=all_docs,
        embedding=embedding,
        persist_directory=PERSIST_DIR,
        collection_metadata={"hnsw:space": "cosine"},
        collection_name="chat_context"
    )

    splits = code_splitter.split_documents(all_docs)
    print(f"🔹 {len(splits)} chunks créés", file=sys.stderr)

    splits = [doc for doc in splits if doc.page_content.strip()]

    # Ajout des statistiques de longueur du split
    bins=50
    counter = Counter(int(len(split.page_content) / bins) * bins for split in splits)
    print("🔹 Histogramme de la longueur des splits :", file=sys.stderr)
    max_value = max(counter.values())
    for length, count in sorted(counter.items()):
        normalized_count = int((count / max_value) * 50)
        print(f"   {length+1:6}-{length+bins:-6}: { '#' * normalized_count}", file=sys.stderr)

    print(f"🔹 {len(splits)} fragments non vides à intégrer", file=sys.stderr)
    # Créer ou recharger Chroma

    # 2. Vectorstore pour la Génération (splits courts)
    gen_collection = Chroma.from_documents(
        documents=splits,  # Morceaux de 200-400 tokens
        embedding=embedding,
        collection_name="code_completion",
        persist_directory=PERSIST_DIR,
        collection_metadata={"hnsw:space": "cosine"}
    )

    vectorstore = {
        "chat": chat_collection,
        "generate": gen_collection
       }

    print("🔹 Vectorstore créé", file=sys.stderr)

# --- Formatage du contexte ---
def format_context(docs: list) -> str:
    context = []
            
    extraits = {}
    print("🔹 Les fichiers suivants ont été selectionnés:", file=sys.stderr)

    for i, doc in enumerate(docs):
        source = doc.metadata.get('source', 'unknown')
        filename = os.path.basename(source)
        print(f"  🔹 {filename} -- extrait {i+1} --", file=sys.stderr)
        if doc.page_content not in extraits:
            extraits[doc.page_content] = True
            context.append(f"### Fichier: {filename} (Extrait {i+1}) ###")
            context.append(textwrap.indent(doc.page_content, '    '))
            print(f"  🔹 {filename} -- fin extrait {i+1} --", file=sys.stderr)
        else:
            print(f"  🔸 {filename} -- extrait {i+1} duppliqué et éliminé --", file=sys.stderr)

    return "\n\n".join(context)

def format_iso_time_with_ns():
    # 1. Obtenir le timestamp actuel avec nanosecondes
    current_time_ns = time.time_ns()
    
    # 2. Convertir en datetime avec timezone locale
    dt = datetime.fromtimestamp(current_time_ns / 1e9).astimezone()
    
    # 3. Formater avec les nanosecondes et décalage horaire
    # - Extraire les nanosecondes
    nanoseconds = current_time_ns % 10**9
    
    # - Formater la partie datetime de base
    base_format = dt.strftime("%Y-%m-%dT%H:%M:%S")
    
    # - Ajouter les nanosecondes (9 chiffres)
    nano_format = f".{nanoseconds:09d}"
    
    # - Formater le décalage horaire
    utc_offset = dt.utcoffset()
    offset_hours = utc_offset.total_seconds() // 3600
    offset_minutes = (utc_offset.total_seconds() % 3600) // 60
    offset_sign = '-' if offset_hours < 0 else '+'
    offset_format = f"{offset_sign}{abs(int(offset_hours)):02d}:{int(offset_minutes):02d}"
    
    return base_format + nano_format + offset_format


# Modèles Pydantic pour l'API compatible Ollama
class GenerateRequest(BaseModel):
    model: str
    prompt: str
    system: Optional[str] = None
    template: Optional[str] = None
    stream: Optional[bool] = False
    options: Optional[Dict[str, Any]] = None

class ChatMessage(BaseModel):
    role: str
    content: str
    images: Optional[List[str]] = None

class ChatRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    format: Optional[str] = None
    options: Optional[Dict[str, Any]] = None
    stream: bool = True
    keep_alive: Optional[Union[str, int]] = None  # Modification ici

class EmbeddingRequest(BaseModel):
    model: str
    prompt: str
    options: Optional[Dict[str, Any]] = None

class EmbeddingResponse(BaseModel):
    embedding: List[float]

# Fonctions utilitaires
async def perform_rag_search(mode: Literal["generate", "chat"], prompt: str, k: int = 4) -> str:
    """Effectue une recherche RAG et retourne le contexte"""
    build_vectorstore()
    
    rag_docs = vectorstore[mode].similarity_search(prompt, k=k)
    return format_context(rag_docs) if rag_docs else "Aucun contexte trouvé."

async def perform_web_search(prompt: str, k: int = 2) -> str:
    """Effectue une recherche web et retourne les résultats"""
    if not DDGS_SEARCH_ENABLED:
        return "Recherche web désactivée"
    
    try:
        from ddgs import DDGS
        with DDGS() as ddgs:
            results = list(ddgs.text(prompt, max_results=k))
            print(f"🔹 {len(results)} résultats trouvés sur le web", file=sys.stderr)
            for i, r in enumerate(results):  
                print(f" - {i+1}. {r['title']} : {r['href']}", file=sys.stderr)
            
            web_info = "\n".join(
                f"- [{r['title']}]({r['href']}): {r['body'][:150]}..." 
                for r in results
            ) if results else "Aucun résultat web trouvé."
    except Exception as e:
        return f"Erreur recherche web: {str(e)}"

def build_enhanced_prompt(
                 mode: Literal["generate", "chat"], 
                 question: str,
                 rag_context: str,
                 web_context: str):
    
    if mode == "chat":
        prompt = f"""
# Consigne

Vous êtes un expert en programmation {PROG_LANG}. Répondez à la question en utilisant le contexte fourni (extraits de code) et les informations web si disponibles.
Autant que possible tu indiqueras tes sources, url, nom du fichier source...

# Contexte de la question:

## **Contexte Code (extraits pertinents):**

{rag_context}

## **Informations Web:**

{web_context}

## **Instructions:**
- Répondez de manière concise et précise à la question
- Si la réponse se trouve dans le contexte code, citez le fichier et l'extrait correspondant
- Si vous utilisez les informations web, citez la source
- Si la question est en anglais, répondez en anglais. Sinon, en français
- Pour les extraits de code, conservez le formatage et l'indentation

# **Question** à laquelle tu dois répondre

{question}
"""
    else:
        prompt=f"""
# Consigne

Vous êtes un expert en programmation {PROG_LANG}. Essayer de concevoir un petit bout de code permetant de résoudre la question

# Contexte de la question:

## **Contexte Code (extraits pertinents):**

{rag_context}

## **Instructions:**

- Rédigez les commentaires de code dans la même langue que le code qui vous est fourni. À défaut en anglais.
- Nommez les variables dans la même langue que le code qui vous est fourni. À défaut en anglais.
- Si la question est en anglais, répondez en anglais. Sinon, en français
- Pour les extraits de code, conservez le formatage et l'indentation

# **Question** à laquelle tu dois répondre

{question}
"""

    return prompt

# --- Fonctions utilitaires factorisées ---
async def build_augmented_prompt(
                 mode: Literal["generate", "chat"], 
                 question: str
        ) -> str:
    """Construit un prompt enrichi avec contextes RAG et web"""

    rag_context = await perform_rag_search(mode,question,k= 2 if mode=='chat' else 8)

    if mode == "chat":
        web_context = await perform_web_search(question)
    else:
        web_context = ""

    return build_enhanced_prompt(
        mode=mode,
        question=question,
        rag_context=rag_context,
        web_context=web_context
    )

async def _stream_ollama_response(response: httpx.Response, model_name: str):
    """Générateur pour le streaming de la réponse de chat au format Ollama"""
    start_time = datetime.now(timezone.utc).isoformat()
    async for line in response.aiter_lines():
        if line.strip():
            try:
                chunk = json.loads(line)
                
                # Construction du message conforme à l'API Ollama
                message_chunk = {
                    "model": model_name,
                    "created_at": start_time,
                    "message": {
                        "role": "assistant",
                        "content": chunk.get("message", {}).get("content", "") if "message" in chunk else chunk.get("content", "")
                    },
                    "done": chunk.get("done", False)
                }
                
                # Ajout des champs optionnels
                for field in ["total_duration", "load_duration", "prompt_eval_count", "eval_count"]:
                    if field in chunk:
                        message_chunk[field] = chunk[field]
                
                yield f"data: {json.dumps(message_chunk)}\n\n"
                
            except json.JSONDecodeError:
                yield "event: error\ndata: Invalid JSON chunk\n\n"
    
    yield "event: end\ndata: Stream completed\n\n"


# --- Endpoints API ---
@app.post("/api/generate")
async def generate_endpoint(request_data: GenerateRequest = Body(...)):
    try:
        # Construction du prompt enrichi
        enhanced_prompt = await build_augmented_prompt("generate",request_data.prompt)
        
        # Appel à Ollama
        ollama_payload = {
            "model": GENERATE_MODEL,
            "prompt": enhanced_prompt,
            "stream": request_data.stream,
            "options": request_data.options or {}
        }
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                f"{OLLAMA_HOST}/api/generate",
                json=ollama_payload,
                timeout=QUERY_TIMEOUT
            )
            response.raise_for_status()
            
            if not request_data.stream:
                return response.json()
            else:
                # --- CORRECTION DU STREAMING ---
                async def generate_stream():
                    """Générateur pour le streaming des résultats"""
                    async for chunk in response.aiter_text():
                        # Transférer directement les chunks
                        yield chunk
                        
                    # Fermeture propre du stream
                    # yield "data: [DONE]\n\n"
                
                # Utilisez text/plain au lieu de text/event-stream
                return StreamingResponse(
                    generate_stream(),
                    media_type="text/plain",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-Accel-Buffering": "no"
                    }
                )
                        
    except httpx.RequestError as e:
        raise HTTPException(500, f"Erreur de connexion à Ollama: {str(e)}")
    except Exception as e:
        raise HTTPException(500, f"Erreur interne: {str(e)}")

@app.post("/api/chat")
async def chat_endpoint(request_data: ChatRequest):
    try:
        messages = [msg.dict() for msg in request_data.messages]
        
        if messages and messages[-1]["role"] == "user":
            original_question = messages[-1]["content"]
            try:
                # Limiter la taille du contexte
                augmented_prompt = await build_augmented_prompt("chat",original_question)
                messages[-1]["content"] = augmented_prompt[-8000:]  # Truncate to the last 8000 chars
                print(f"🔹 Prompt enrichi ({len(augmented_prompt)} caractères)", file=sys.stderr)
            except Exception as e:
                print(f"⚠️ Erreur d'enrichissement: {str(e)}", file=sys.stderr)
                messages[-1]["content"] = original_question  # Fallback to original
        
        # Préparation du payload pour Ollama
        ollama_payload = {
            "model": CHAT_MODEL,
            "messages": messages,
            "stream": request_data.stream,
            "options": request_data.options or {}
        }

        if request_data.keep_alive is not None:
            if isinstance(request_data.keep_alive, int):
                ollama_payload["keep_alive"] = f"{request_data.keep_alive}s"
            else:
                ollama_payload["keep_alive"] = request_data.keep_alive
        
        sopload = json.dumps(ollama_payload, indent=2)
        print(f" 🔹 Taille du Payload vers Ollama : {len(sopload)} octets...", file=sys.stderr)
        print(f" 🔹 Début du payload : {sopload}...", file=sys.stderr)

        # Appel à Ollama
        async with httpx.AsyncClient() as client:
            try:
                response = await client.post(
                    f"{OLLAMA_HOST}/api/chat",
                    json=ollama_payload,
                    timeout=QUERY_TIMEOUT
                )
                response.raise_for_status()
                
                # Gestion des réponses NON-STREAMING
                if not request_data.stream:
                    return response.json()
                
                # --- CORRECTION DU STREAMING ---
                async def generate_stream():
                    """Générateur pour le streaming des résultats"""
                    async for chunk in response.aiter_text():
                        # Transférer directement les chunks
                        yield chunk
                        
                    # Fermeture propre du stream
                    # yield "data: [DONE]\n\n"
                
                # Utilisez text/plain au lieu de text/event-stream
                return StreamingResponse(
                    generate_stream(),
                    media_type="text/plain",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "X-Accel-Buffering": "no"
                    }
                )
    # ... [gestion des erreurs existante] ...
            except httpx.HTTPStatusError as e:
                error_detail = e.response.text if e.response else str(e)
                print(f"🚨 Erreur HTTP Ollama ({e.response.status_code}): {error_detail}", file=sys.stderr)
                raise HTTPException(502, f"Erreur Ollama: {error_detail}")
            except httpx.RequestError as e:
                print(f"🚨 Erreur réseau Ollama: {str(e)}", file=sys.stderr)
                raise HTTPException(503, f"Ollama non disponible: {str(e)}")
                
    except Exception as e:
        print(f"🚨 Erreur interne: {traceback.format_exc()}", file=sys.stderr)
        raise HTTPException(500, f"Erreur interne: {str(e)}")

@app.post("/api/embeddings")
async def embeddings(request: EmbeddingRequest):
    """Proxy direct pour les embeddings"""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                f"{OLLAMA_HOST}/api/embeddings",
                json=request.dict()
            )
            response.raise_for_status()
            return response.json()
            
        except httpx.RequestError as e:
            raise HTTPException(500, f"Erreur de connexion à Ollama: {str(e)}")

@app.get("/api/tags")
async def list_models():
    """Proxy pour lister les modèles disponibles"""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(f"{OLLAMA_HOST}/api/tags")
            response.raise_for_status()
            return response.json()
            
        except httpx.RequestError as e:
            raise HTTPException(500, f"Erreur de connexion à Ollama: {str(e)}")

# Initialisation du vectorstore (à adapter à votre code)
@app.on_event("startup")
async def startup_event():
    global vectorstore
    build_vectorstore()
    print("🔹 Initialisation du serveur proxy Ollama+RAG")

# Endpoint supplémentaire pour le contrôle
@app.get("/control/enable_web_search")
async def enable_web_search(enabled: bool = True):
    global DDGS_SEARCH_ENABLED
    DDGS_SEARCH_ENABLED = enabled
    return {"status": "success", "web_search_enabled": enabled}

# Endpoint de debug simplifié
@app.post("/debug")
async def debug_endpoint(request: Request):
    """Endpoint de débogage simplifié"""
    try:
        body = await request.json()
        return {
            "status": "success",
            "received_body": body
        }
    except json.JSONDecodeError:
        raise HTTPException(400, "Invalid JSON format")