"""Real CLIP inference behind the OpenAI-compatible embedding contract."""
import base64,io,logging
import torch
from PIL import Image
from transformers import CLIPModel,CLIPProcessor
from fastapi import FastAPI
model_id='openai/clip-vit-base-patch32'
model=CLIPModel.from_pretrained(model_id)
processor=CLIPProcessor.from_pretrained(model_id)
model.eval()
app=FastAPI()
@app.post('/v1/embeddings')
def embeddings(body:dict):
    if 'messages' in body:
        items=[c for m in body['messages'] if m['role']=='user' for c in m['content']]
        images=[Image.open(io.BytesIO(base64.b64decode(c['image_url']['url'].split(',',1)[1]))).convert('RGB') for c in items if c['type']=='image_url']
        with torch.inference_mode():
            result=model.get_image_features(**processor(images=images,return_tensors='pt'))
        kind='image'
    else:
        texts=body['input'];texts=[texts] if isinstance(texts,str) else texts
        with torch.inference_mode():
            result=model.get_text_features(**processor(text=texts,return_tensors='pt',padding=True,truncation=True,max_length=77))
        kind='text'
    if not isinstance(result,torch.Tensor):result=result.pooler_output
    result=result/result.norm(dim=-1,keepdim=True)
    logging.warning('Real CLIP inference kind=%s count=%d dimensions=%d',kind,result.shape[0],result.shape[1])
    return {'object':'list','model':model_id,'data':[{'object':'embedding','index':i,'embedding':v} for i,v in enumerate(result.tolist())],'usage':{'prompt_tokens':1,'total_tokens':1}}

@app.post('/v1/rerank')
def rerank(body:dict):
    def encode(value):
        if isinstance(value,str):return embeddings({'input':[value]})['data'][0]['embedding']
        content=value['content']
        images=[c for c in content if c['type']=='image_url']
        if images:return embeddings({'messages':[{'role':'user','content':images}]})['data'][0]['embedding']
        return encode(' '.join(c.get('text','') for c in content))
    query=torch.tensor(encode(body['query']))
    docs=body['documents']
    results=[{'index':i,'relevance_score':float(torch.dot(query,torch.tensor(encode(d)))), 'document':{'text':d if isinstance(d,str) else 'image'}} for i,d in enumerate(docs)]
    logging.warning('REAL_CLIP_RERANK query_type=%s document_types=%s',type(body['query']).__name__,[type(d).__name__ for d in docs])
    return {'results':sorted(results,key=lambda x:x['relevance_score'],reverse=True)[:body.get('top_n',len(docs))]}

"""Create and replay a disposable Dify fixture through authenticated HTTP APIs."""
import argparse
import base64
import io
import json
import os
import time
from pathlib import Path

import requests
from PIL import Image

parser = argparse.ArgumentParser()
parser.add_argument('mode', choices=['setup', 'before', 'after'])
parser.add_argument('--state', default='fixture.json')
args = parser.parse_args()
base = os.environ.get('DIFY_URL', 'http://127.0.0.1:15011').rstrip('/')
endpoint = os.environ.get('MODEL_URL', 'http://host.docker.internal:15013/v1')
provider = 'langgenius/openai_api_compatible/openai_api_compatible'
session = requests.Session()
session.trust_env = False

def call(method, path, body=None, files=None, service=False):
    headers = {}
    if service:
        headers['Authorization'] = 'Bearer ' + os.environ['DIFY_DATASET_API_KEY']
    else:
        for cookie in session.cookies:
            if 'csrf' in cookie.name:
                headers['X-CSRF-Token'] = cookie.value
    response = session.request(method, base + path, json=body, files=files, headers=headers, timeout=120)
    if not response.ok:
        raise RuntimeError(f"{response.status_code} {path}: {response.text}")
    return response.json()

if args.mode == 'setup':
    call('POST', '/console/api/login', {
        'email': os.environ['DIFY_EMAIL'],
        'password': base64.b64encode(os.environ['DIFY_PASSWORD'].encode()).decode(),
        'language': 'en-US', 'remember_me': False,
    })
    suffix = str(time.time_ns())
    embedding = 'clip-embedding-' + suffix
    reranker = 'clip-reranker-' + suffix
    for name, kind in [(embedding, 'text-embedding'), (reranker, 'rerank')]:
        call('POST', '/console/api/workspaces/current/model-providers/' + provider + '/models/credentials', {
            'model': name, 'model_type': kind, 'name': 'Reproduction',
            'credentials': {'endpoint_url': endpoint, 'api_key': 'local-test',
                            'vision_support': 'support', 'context_size': '77', 'max_chunks': '16'},
        })
    call('POST', '/console/api/workspaces/current/default-model', {
        'model_settings': [{'model_type': 'text-embedding', 'provider': provider, 'model': embedding}],
    })
    dataset = call('POST', '/console/api/datasets', {
        'name': 'Vision rerank repro ' + suffix, 'indexing_technique': 'high_quality',
    })['id']
    call('PATCH', '/console/api/datasets/' + dataset, {
        'indexing_technique': 'high_quality', 'embedding_model': embedding,
        'embedding_model_provider': provider, 'is_multimodal': True,
        'retrieval_model': {'search_method': 'semantic_search', 'reranking_enable': False,
                            'top_k': 10, 'score_threshold_enabled': False},
    })
    document = call('POST', f'/v1/datasets/{dataset}/document/create-by-text', {
        'name': 'Red square', 'text': 'A red square diagram.', 'indexing_technique': 'high_quality',
        'embedding_model': embedding, 'embedding_model_provider': provider,
        'process_rule': {'mode': 'automatic'},
    }, service=True)['document']['id']
    for _ in range(120):
        segments = call('GET', f'/console/api/datasets/{dataset}/documents/{document}/segments')['data']
        if segments and all(s['status'] == 'completed' for s in segments):
            break
        time.sleep(1)
    else:
        raise RuntimeError('Indexing did not finish within 120 seconds')
    png = io.BytesIO()
    Image.new('RGB', (224, 224), (255, 0, 0)).save(png, format='PNG')
    upload = call('POST', '/console/api/files/upload', files={'file': ('red-square.png', png.getvalue(), 'image/png')})['id']
    for label in ['A', 'B']:
        call('POST', f'/console/api/datasets/{dataset}/documents/{document}/segment', {
            'content': f'Segment {label}: a red square diagram.', 'attachment_ids': [upload],
        })
    Path(args.state).write_text(json.dumps({'dataset': dataset, 'reranker': reranker}))
    print('Created fixture: three text segments; A and B share one PNG. Run before/after after indexing.')
else:
    state = json.loads(Path(args.state).read_text())
    evidence = []
    for enabled in [False, True]:
        body = {'query': 'a red square', 'retrieval_model': {
            'search_method': 'semantic_search', 'reranking_enable': enabled,
            'reranking_model': {'reranking_provider_name': provider, 'reranking_model_name': state['reranker']},
            'top_k': 10, 'score_threshold_enabled': False,
        }}
        result = call('POST', '/v1/datasets/' + state['dataset'] + '/retrieve', body, service=True)
        image_ids = sorted({f['id'] for r in result['records'] for f in r.get('files', [])})
        row = {'reranking_enable': enabled, 'status': 200, 'image_ids': image_ids,
               'records': [{'segment_id': r['segment']['id'], 'score': r['score'],
                            'file_ids': [f['id'] for f in r.get('files', [])]} for r in result['records']]}
        evidence.append(row)
        print(f'reranking_enable={enabled}: HTTP 200, returned images={len(image_ids)}')
    Path(args.mode + '.json').write_text(json.dumps(evidence, indent=2))
    assert [len(e['image_ids']) for e in evidence] == ([1, 0] if args.mode == 'before' else [1, 1])