import hashlib, json, time
from datetime import datetime, timezone
from typing import Any
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.api.deps import current_user
from app.core.database import get_db
from app.core.security import encrypt_secret
from app.models.entities import *
from app.services.ai_rag import chunk_text, local_embedding, retrieve, redact, resolve_policy, render_prompt, call_provider, accessible_kb

router=APIRouter(prefix='/api/v1/v008',tags=['V008 Enterprise AI'])

class ProviderPayload(BaseModel):
    name:str; provider_type:str; base_url:str=''; model_name:str=''; config:dict={}; is_local:bool=False; enabled:bool=True
class PolicyPayload(BaseModel):
    name:str; department_id:str|None=None; allowed_provider_ids:list[str]=[]; allowed_roles:list[str]=[]; max_context_chars:int=30000; max_output_tokens:int=1200; allow_external:bool=False; require_approved_prompt:bool=True; redact_patterns:list[str]=[]; system_rules:str=''; enabled:bool=True
class KBPayload(BaseModel):
    name:str; description:str=''; department_id:str|None=None; classification:str='internal'; allowed_roles:list[str]=[]; enabled:bool=True
class DocumentPayload(BaseModel):
    title:str; content:str=Field(min_length=1,max_length=2_000_000); source_type:str='text'; source_uri:str=''; metadata:dict={}
class PromptPayload(BaseModel):
    name:str; description:str=''; template:str=Field(min_length=1,max_length=100000); variables:list[str]=[]; department_id:str|None=None; status:str='draft'
class ConversationPayload(BaseModel):
    title:str='Nova conversa'; provider_id:str; knowledge_base_ids:list[str]=[]; prompt_template_id:str|None=None
class AskPayload(BaseModel):
    question:str=Field(min_length=1,max_length=20000); variables:dict={}; top_k:int=6

def elevated(u): return u.role in {Role.admin,Role.director,Role.data_steward}
def own_or_dept(obj,u): return elevated(u) or getattr(obj,'owner_id',None)==u.id or (getattr(obj,'department_id',None) and getattr(obj,'department_id',None)==u.department_id)

@router.post('/providers')
def create_provider(p:ProviderPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    if not elevated(u): raise HTTPException(403,'Somente governança pode cadastrar provedores')
    if p.provider_type not in {'openai','gemini','openai_compatible','local','mock'}: raise HTTPException(400,'Tipo de provedor não suportado')
    row=EnterpriseAIProvider(name=p.name,provider_type=p.provider_type,base_url=p.base_url,model_name=p.model_name,encrypted_config=encrypt_secret(json.dumps(p.config)),is_local=p.is_local,enabled=p.enabled,created_by=u.id)
    db.add(row); db.commit(); db.refresh(row); return {'id':row.id}

@router.get('/providers')
def list_providers(db:Session=Depends(get_db),u:User=Depends(current_user)):
    policy=resolve_policy(db,u); allowed=set(policy.allowed_provider_ids or []) if policy else set()
    rows=db.scalars(select(EnterpriseAIProvider).where(EnterpriseAIProvider.enabled==True).order_by(EnterpriseAIProvider.name)).all()
    return [{'id':x.id,'name':x.name,'provider_type':x.provider_type,'model_name':x.model_name,'is_local':x.is_local,'enabled':x.enabled} for x in rows if not allowed or x.id in allowed]

@router.post('/policies')
def create_policy(p:PolicyPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    if not elevated(u): raise HTTPException(403,'Somente governança pode cadastrar políticas')
    row=AIPolicy(**p.model_dump(),created_by=u.id); db.add(row); db.commit(); db.refresh(row); return {'id':row.id}

@router.get('/policy/current')
def current_policy(db:Session=Depends(get_db),u:User=Depends(current_user)):
    p=resolve_policy(db,u)
    if not p: return {'configured':False}
    return {'configured':True,'id':p.id,'name':p.name,'allow_external':p.allow_external,'require_approved_prompt':p.require_approved_prompt,'max_context_chars':p.max_context_chars,'max_output_tokens':p.max_output_tokens}

@router.post('/knowledge-bases')
def create_kb(p:KBPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    dep=p.department_id or u.department_id
    if p.department_id and p.department_id!=u.department_id and not elevated(u): raise HTTPException(403,'Departamento não autorizado')
    row=KnowledgeBase(**p.model_dump(exclude={'department_id'}),department_id=dep,owner_id=u.id); db.add(row); db.commit(); db.refresh(row); return {'id':row.id}

@router.get('/knowledge-bases')
def list_kbs(db:Session=Depends(get_db),u:User=Depends(current_user)):
    rows=db.scalars(select(KnowledgeBase).where(KnowledgeBase.enabled==True).order_by(KnowledgeBase.created_at.desc())).all()
    return [{'id':x.id,'name':x.name,'description':x.description,'classification':x.classification,'department_id':x.department_id} for x in rows if accessible_kb(x,u)]

@router.post('/knowledge-bases/{kb_id}/documents')
def add_document(kb_id:str,p:DocumentPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    kb=db.get(KnowledgeBase,kb_id)
    if not kb or not own_or_dept(kb,u): raise HTTPException(404,'Base de conhecimento não encontrada')
    digest=hashlib.sha256(p.content.encode()).hexdigest()
    existing=db.scalar(select(KnowledgeDocument).where(KnowledgeDocument.knowledge_base_id==kb_id,KnowledgeDocument.content_hash==digest))
    if existing: return {'id':existing.id,'deduplicated':True,'chunks':db.query(KnowledgeChunk).filter_by(document_id=existing.id).count()}
    doc=KnowledgeDocument(knowledge_base_id=kb_id,title=p.title,source_type=p.source_type,source_uri=p.source_uri,content_hash=digest,metadata_json=p.metadata,created_by=u.id)
    db.add(doc); db.flush(); pieces=chunk_text(p.content)
    for idx,text in enumerate(pieces): db.add(KnowledgeChunk(document_id=doc.id,chunk_index=idx,content=text,embedding=local_embedding(text),token_estimate=max(1,len(text)//4),metadata_json=p.metadata))
    db.commit(); return {'id':doc.id,'deduplicated':False,'chunks':len(pieces)}

@router.get('/knowledge-bases/{kb_id}/documents')
def list_documents(kb_id:str,db:Session=Depends(get_db),u:User=Depends(current_user)):
    kb=db.get(KnowledgeBase,kb_id)
    if not kb or not accessible_kb(kb,u): raise HTTPException(404,'Base não encontrada')
    rows=db.scalars(select(KnowledgeDocument).where(KnowledgeDocument.knowledge_base_id==kb_id).order_by(KnowledgeDocument.created_at.desc())).all()
    return [{'id':x.id,'title':x.title,'source_type':x.source_type,'status':x.status,'created_at':x.created_at,'metadata':x.metadata_json} for x in rows]

@router.post('/prompts')
def create_prompt(p:PromptPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    if p.status=='approved' and not elevated(u): raise HTTPException(403,'Aprovação restrita')
    row=PromptTemplate(**p.model_dump(),owner_id=u.id,approved_by=u.id if p.status=='approved' else None,approved_at=datetime.now(timezone.utc) if p.status=='approved' else None)
    db.add(row); db.commit(); db.refresh(row); return {'id':row.id,'version':row.version}

@router.get('/prompts')
def list_prompts(db:Session=Depends(get_db),u:User=Depends(current_user)):
    q=select(PromptTemplate).where((PromptTemplate.owner_id==u.id)|(PromptTemplate.department_id==u.department_id)|(PromptTemplate.status=='approved')).order_by(PromptTemplate.created_at.desc())
    return [{'id':x.id,'name':x.name,'description':x.description,'template':x.template,'variables':x.variables,'status':x.status,'version':x.version} for x in db.scalars(q).all()]

@router.post('/prompts/{pid}/approve')
def approve_prompt(pid:str,db:Session=Depends(get_db),u:User=Depends(current_user)):
    if not elevated(u): raise HTTPException(403,'Aprovação restrita')
    row=db.get(PromptTemplate,pid)
    if not row: raise HTTPException(404,'Prompt não encontrado')
    row.status='approved'; row.approved_by=u.id; row.approved_at=datetime.now(timezone.utc); db.commit(); return {'status':'approved'}

@router.post('/conversations')
def create_conversation(p:ConversationPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    provider=db.get(EnterpriseAIProvider,p.provider_id); policy=resolve_policy(db,u)
    if not provider or not provider.enabled: raise HTTPException(404,'Provedor não encontrado')
    if policy and policy.allowed_provider_ids and provider.id not in policy.allowed_provider_ids: raise HTTPException(403,'Provedor não autorizado pela política')
    if provider.is_local is False and policy and not policy.allow_external: raise HTTPException(403,'Transferência externa bloqueada')
    for kb_id in p.knowledge_base_ids:
        kb=db.get(KnowledgeBase,kb_id)
        if not kb or not accessible_kb(kb,u): raise HTTPException(403,'Base não autorizada')
    prompt=db.get(PromptTemplate,p.prompt_template_id) if p.prompt_template_id else None
    if policy and policy.require_approved_prompt and (not prompt or prompt.status!='approved'): raise HTTPException(403,'Prompt aprovado obrigatório')
    row=EnterpriseAIConversation(**p.model_dump(),user_id=u.id,department_id=u.department_id); db.add(row); db.commit(); db.refresh(row); return {'id':row.id}

@router.get('/conversations')
def list_conversations(db:Session=Depends(get_db),u:User=Depends(current_user)):
    rows=db.scalars(select(EnterpriseAIConversation).where(EnterpriseAIConversation.user_id==u.id).order_by(EnterpriseAIConversation.updated_at.desc())).all()
    return [{'id':x.id,'title':x.title,'provider_id':x.provider_id,'knowledge_base_ids':x.knowledge_base_ids,'created_at':x.created_at,'updated_at':x.updated_at} for x in rows]

@router.get('/conversations/{cid}/messages')
def list_messages(cid:str,db:Session=Depends(get_db),u:User=Depends(current_user)):
    conv=db.get(EnterpriseAIConversation,cid)
    if not conv or conv.user_id!=u.id: raise HTTPException(404,'Conversa não encontrada')
    rows=db.scalars(select(AIMessage).where(AIMessage.conversation_id==cid).order_by(AIMessage.created_at)).all()
    return [{'id':x.id,'role':x.role,'content':x.content,'citations':x.citations,'safety_flags':x.safety_flags,'created_at':x.created_at} for x in rows]

@router.post('/conversations/{cid}/ask')
def ask(cid:str,p:AskPayload,db:Session=Depends(get_db),u:User=Depends(current_user)):
    conv=db.get(EnterpriseAIConversation,cid)
    if not conv or conv.user_id!=u.id: raise HTTPException(404,'Conversa não encontrada')
    provider=db.get(EnterpriseAIProvider,conv.provider_id); policy=resolve_policy(db,u)
    if not provider or not provider.enabled: raise HTTPException(409,'Provedor indisponível')
    if policy and policy.allowed_provider_ids and provider.id not in policy.allowed_provider_ids: raise HTTPException(403,'Provedor bloqueado')
    external=not provider.is_local and provider.provider_type!='mock'
    if external and policy and not policy.allow_external: raise HTTPException(403,'Transferência externa bloqueada')
    prompt_template=db.get(PromptTemplate,conv.prompt_template_id) if conv.prompt_template_id else None
    if policy and policy.require_approved_prompt and (not prompt_template or prompt_template.status!='approved'): raise HTTPException(403,'Prompt aprovado obrigatório')
    question,flags=redact(p.question,policy.redact_patterns if policy else [])
    hits=retrieve(db,u,conv.knowledge_base_ids,question,p.top_k)
    max_chars=policy.max_context_chars if policy else 30000; context=''; citations=[]
    for h in hits:
        block=f"\n[FONTE {len(citations)+1}: {h['title']}]\n{h['content']}\n"
        if len(context)+len(block)>max_chars: break
        context+=block; citations.append({'document_id':h['document_id'],'chunk_id':h['chunk_id'],'title':h['title'],'score':round(h['score'],4)})
    user_prompt=render_prompt(prompt_template,question,p.variables)+"\n\nCONTEXTO AUTORIZADO:\n"+context
    user_prompt,more_flags=redact(user_prompt,policy.redact_patterns if policy else []); flags.extend(more_flags)
    system=(policy.system_rules if policy else '')+"\nResponda somente com base no contexto autorizado. Quando não houver evidência suficiente, declare isso. Cite as fontes pelo número."
    db.add(AIMessage(conversation_id=cid,role='user',content=p.question,safety_flags=flags)); db.flush(); started=time.monotonic()
    try:
        answer,usage=call_provider(provider,system,user_prompt,policy.max_output_tokens if policy else 1200)
        latency=int((time.monotonic()-started)*1000)
        msg=AIMessage(conversation_id=cid,role='assistant',content=answer,citations=citations,safety_flags=flags,prompt_tokens=int(usage.get('prompt_tokens',usage.get('promptTokenCount',0)) or 0),completion_tokens=int(usage.get('completion_tokens',usage.get('candidatesTokenCount',0)) or 0),latency_ms=latency)
        db.add(msg); db.add(AIAuditEvent(user_id=u.id,conversation_id=cid,action='rag_query',provider_id=provider.id,external_transfer=external,context_chars=len(context),details={'knowledge_bases':conv.knowledge_base_ids,'citations':len(citations),'safety_flags':flags})); conv.updated_at=datetime.now(timezone.utc); db.commit(); db.refresh(msg)
        return {'message_id':msg.id,'answer':answer,'citations':citations,'safety_flags':flags,'latency_ms':latency}
    except Exception as exc:
        db.rollback(); db.add(AIAuditEvent(user_id=u.id,conversation_id=cid,action='rag_query_failed',provider_id=provider.id,external_transfer=external,context_chars=len(context),details={'error':str(exc)[:1000]})); db.commit(); raise HTTPException(502,'Falha controlada ao consultar o provedor')

@router.get('/audit')
def audit(db:Session=Depends(get_db),u:User=Depends(current_user)):
    if not elevated(u) and u.role!=Role.auditor: raise HTTPException(403,'Acesso restrito')
    q=select(AIAuditEvent).order_by(AIAuditEvent.created_at.desc()).limit(500)
    return [{'id':x.id,'user_id':x.user_id,'conversation_id':x.conversation_id,'action':x.action,'provider_id':x.provider_id,'external_transfer':x.external_transfer,'context_chars':x.context_chars,'details':x.details,'created_at':x.created_at} for x in db.scalars(q).all()]
