import os, re, time, math, unicodedata, logging
from contextlib import asynccontextmanager
from collections import Counter
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import numpy as np
import torch
import torch.nn as nn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM
from keybert import KeyBERT
from huggingface_hub import hf_hub_download, snapshot_download

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("liftup")

HF_USERNAME = os.getenv("HF_USERNAME", "Engin34")
HF_TOKEN    = os.getenv("HF_TOKEN", "")
GEMINI_KEY  = os.getenv("GEMINI_API_KEY", "")

bert_model   = None
bert_tok     = None
kw_model     = None
generator    = None
clf          = None
TOP_KEYWORDS = None
_ext_model   = None
_EXT_OK      = False
ALL_CATS     = []


# ── Taksonomi parser ─────────────────────────────────────────────────

def _temizle(k):
    k = k.replace('\u200b', '').replace('\ufeff', '')
    k = unicodedata.normalize('NFKC', k)
    return re.sub(r'\s+', ' ', k).strip().lower()

def _parantez_ayir(k):
    m = re.match(r'^(.+?)\s*\((.+?)\)\s*$', k)
    if not m:
        return [k]
    ana, ic = m.group(1).strip(), m.group(2).strip()
    if any(a in ic.lower() for a in ['baglam', 'proses', 'analiz', 'anahtarlari']):
        return [ana]
    return [ana] + [p.strip() for p in ic.split('/') if p.strip()]

def _virgul_ayir(metin):
    sonuc, buf, d = [], [], 0
    for c in metin:
        if c == '(':
            d += 1
            buf.append(c)
        elif c == ')':
            d -= 1
            buf.append(c)
        elif c == ',' and d == 0:
            sonuc.append(''.join(buf))
            buf = []
        else:
            buf.append(c)
    if buf:
        sonuc.append(''.join(buf))
    return sonuc

def parse_taksonomi(icerik: str) -> Dict:
    icerik = icerik.replace('\u200b', '')
    matches = list(re.finditer(r'^\s*(\d+)\)\s+(.+?)\s*$', icerik, re.MULTILINE))
    tax = {}
    for i, m in enumerate(matches):
        kat   = m.group(2).strip()
        govde = icerik[m.end():(matches[i+1].start() if i+1 < len(matches) else len(icerik))].strip()
        pm    = re.search(r'\((.+)\)', govde, re.DOTALL)
        if not pm:
            continue
        kw_set = set()
        for parca in _virgul_ayir(pm.group(1)):
            for alt in _parantez_ayir(parca.strip()):
                for k in re.split(r'[/]', alt):
                    temiz = _temizle(k)
                    if len(temiz) >= 2:
                        kw_set.add(temiz)
        tax[kat] = {'keywords': kw_set}
    return tax


# ── Taksonomi genisletme sozlugu ─────────────────────────────────────

_EKSTRA_KW = {
    'Yapay Zeka & Veri Bilimi': {
        'derin ogrenme', 'sinir agi', 'transformer', 'lstm', 'gru',
        'makine ogrenmesi', 'yapay sinir', 'classification', 'regresyon',
        'veri madenciligi', 'takviyeli ogrenme', 'xgboost',
    },
    'Yapilar, Mukavemet & Analiz': {
        'sonlu elemanlar', 'fea', 'burkulma', 'yorulma', 'hasar',
        'gerilme', 'deformasyon', 'mod analizi', 'titresim', 'nastran',
        'abaqus', 'ansys', 'mukavemet',
    },
    'Goruntu Isleme & Algilama': {
        'nesne tespiti', 'segmentasyon', 'lidar', 'radar', 'yolo',
        'resnet', 'konvolusyon', 'termal', 'infrared', 'optik akis',
    },
    'Otonom Sistemler & Robotik': {
        'otonom', 'iha', 'uav', 'suru', 'navigasyon', 'haritalama',
        'slam', 'path planning', 'insansiz',
    },
    'Ucus Mekanigi': {
        'stabilite', 'trim', 'manevra', 'irtifa', 'gudum',
        'ucus yorungesi', 'agirlik merkezi', 'balistik',
        'ucus dinamigi', 'manevrabilite', 'yukleme', 'flaperon',
        'elevator', 'rudder', 'aileron', 'moment', 'kaldirma',
        'surukleme', 'aoa', 'hucum acisi', 'dalis', 'tirmanma',
    },
    'Kontrol & Aviyonik': {
        'pid', 'kontrol sistemi', 'geri besleme', 'servo', 'gyro',
        'ivmeolcer', 'imu', 'kalman', 'fpga', 'gomulu',
        'mikrodenetleyici', 'aviyonik', 'ucus bilgisayari',
        'otopilot', 'fms', 'actuator', 'sensor fuzyonu',
    },
    'Eklemeli Imalat': {
        '3d baski', 'slm', 'fdm', 'sls', 'dmls', 'ebm',
        'topoloji optimizasyonu', 'kafes yapi', 'metal baski',
        'polimer', 'eklemeli', 'additive', 'katman katman',
        'tarama stratejisi', 'goz enek', 'yogunluk',
    },
    'Yazilim Gelistirme': {
        'gomulu yazilim', 'rtos', 'misra', 'do-178', 'arinc',
        'mil-std', 'yazilim dogrulama', 'test otomasyonu',
        'model tabanli tasarim', 'simulink', 'model based',
    },
}


# ── Veri siniflari ───────────────────────────────────────────────────

@dataclass
class KategoriSkoru:
    kategori:      str
    final_skor:    float
    keyword_skor:  float
    semantic_skor: float
    eslesmeler:    list = field(default_factory=list)


# ── Hibrit siniflandirici ────────────────────────────────────────────

class HibritSiniflandirici:

    def __init__(self, taxonomy, embedder,
                 keyword_weight=0.35, semantic_weight=0.65, title_boost=2.0):
        self.taxonomy    = self._expand(
            {c: {'keywords': {str(k).lower().strip()
                              for k in d.get('keywords', set()) if str(k).strip()}}
             for c, d in taxonomy.items()}
        )
        self.kw_w        = keyword_weight
        self.sem_w       = semantic_weight
        self.title_boost = title_boost
        self.embedder    = embedder
        self.idf         = self._idf()
        self.centroids   = self._centroids()
        log.info("Hibrit siniflandirici hazir: %d kategori", len(self.taxonomy))

    def _expand(self, tax):
        import copy
        t = copy.deepcopy(tax)
        for cat, kws in _EKSTRA_KW.items():
            if cat in t:
                t[cat]['keywords'] |= kws
        return t

    def _idf(self):
        cnt = Counter()
        N   = len(self.taxonomy)
        for d in self.taxonomy.values():
            for k in d['keywords']:
                cnt[k] += 1
        return {k: math.log((N + 1) / (v + 1)) + 1 for k, v in cnt.items()}

    def _centroids(self):
        c = {}
        for cat, d in self.taxonomy.items():
            kws = list(d['keywords'])[:30]
            if not kws:
                c[cat] = None
                continue
            embs = self.embedder.encode(kws, show_progress_bar=False, convert_to_numpy=True)
            v    = np.mean(embs, axis=0)
            n    = np.linalg.norm(v)
            c[cat] = v / n if n > 0 else v
        return c

    def _kw_score(self, extracted, title_kws=None):
        ext     = [k.lower().strip() for k in extracted if k and str(k).strip()]
        max_idf = max(self.idf.values(), default=1.0)
        results = {}
        for cat, d in self.taxonomy.items():
            cat_kws      = d['keywords']
            score, eslm  = 0.0, []
            for kw in ext:
                boost = self.title_boost if (title_kws and kw in title_kws) else 1.0
                idf_w = self.idf.get(kw, 1.0)
                if kw in cat_kws:
                    p = 2.0 * boost * idf_w
                    score += p
                    eslm.append(kw)
                elif len(kw) >= 4 and any(kw in ck or ck in kw for ck in cat_kws):
                    p = 1.0 * boost * idf_w
                    score += p
                    eslm.append(kw)
            max_p = max(len(ext) * 2.0 * max_idf, 1e-6)
            results[cat] = (min(score / max_p, 1.0), eslm)
        return results

    def _sem_score(self, extracted, text=None):
        parts = []
        if text and str(text).strip():
            parts.append(str(text).strip())
        if extracted:
            parts.append(' '.join(str(k) for k in extracted))
        if not parts:
            return {c: 0.0 for c in self.taxonomy}
        emb = self.embedder.encode(
            [' | '.join(parts)], show_progress_bar=False, convert_to_numpy=True
        )[0]
        n = np.linalg.norm(emb)
        if n > 0:
            emb = emb / n
        return {
            cat: max(0.0, min(1.0, (float(np.dot(emb, cn)) + 1.0) / 2.0))
            if cn is not None else 0.0
            for cat, cn in self.centroids.items()
        }

    def classify(self, keywords, text=None, title=None, top_k=3):
        title_kws = {w.lower() for w in str(title).split()} if title else None
        kw_r      = self._kw_score(keywords, title_kws)
        sem_s     = self._sem_score(keywords, text)
        scores    = {}
        for c in self.taxonomy:
            kwn, esl  = kw_r[c]
            f         = self.kw_w * kwn + self.sem_w * sem_s[c]
            scores[c] = KategoriSkoru(c, f, kwn, sem_s[c], esl)
        srt = sorted(scores.values(), key=lambda x: x.final_skor, reverse=True)
        return {
            'prediction': srt[0].kategori,
            'confidence': srt[0].final_skor,
            'top_k':      srt[:top_k],
            'all_scores': scores,
        }


# ── BERT modeli ──────────────────────────────────────────────────────

class LiftUpBertModel(nn.Module):

    def __init__(self, num_labels=128):
        super().__init__()
        self.bert       = AutoModel.from_pretrained('dbmdz/bert-base-turkish-cased')
        self.dropout    = nn.Dropout(0.3)
        self.classifier = nn.Linear(768, num_labels)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        return self.classifier(self.dropout(out.last_hidden_state[:, 0]))


# ── Post-processor ───────────────────────────────────────────────────

class SoftPostProcessor:

    BLACKLIST = {'kombinatur', 'hesonomik', 'modulasyonlari', 'difuzorlue', 'optimizasyonlarini'}
    ACRONYMS  = {'CFD', 'FEA', 'CAD', 'ROS', 'CNN', 'AI', 'ML', 'DL', 'IoT', 'GPU', 'SSD', 'API'}

    def fix_case(self, kw):
        parts = []
        for w in kw.split():
            if w.upper() in self.ACRONYMS:
                parts.append(w.upper())
            elif not parts:
                parts.append(w.capitalize())
            else:
                parts.append(w.lower())
        return ' '.join(parts)

    def ok(self, kw):
        return (
            kw.lower() not in self.BLACKLIST
            and 3 <= len(kw) <= 80
            and not re.search(r'[^a-zA-Z\xc7\xe7\u011e\u011f\u0130\u0131\xd6\xf6\u015e\u015f\xdc\xfc\s\-]', kw)
        )

    def process(self, kws, min_kw=3):
        out, seen = [], set()
        for k in kws:
            if not self.ok(k):
                continue
            f = self.fix_case(k)
            if f.lower() not in seen:
                out.append(f)
                seen.add(f.lower())
        return out[:8] if out else kws[:min_kw]


# ── Ek destek modulu ─────────────────────────────────────────────────

def _init_ext(api_key, cat_list):
    global _ext_model, _EXT_OK
    if not api_key:
        return
    try:
        import google.generativeai as genai
        genai.configure(api_key=api_key)
        _ext_model = genai.GenerativeModel('gemini-2.5-flash')
        _ext_model.generate_content('Hi')
        _EXT_OK = True
        log.info('Ek destek modulu hazir')
    except Exception as e:
        _EXT_OK = False
        log.warning('Ek destek modulu: %s', e)


def _ext_keywords(title, abstract):
    if not _EXT_OK or _ext_model is None:
        return []
    try:
        prompt = (
            f'Havacilk ve savunma sanayi alaninda uzman bir arastirmacisin.\n'
            f'Asagidaki projenin teknik anahtar kelimelerini belirle.\n\n'
            f'PROJE BASLIGI: {title}\n'
            f'PROJE OZETI: {abstract[:500]}\n\n'
            f'KURALLAR:\n'
            f'- Yalnizca bu projede gecen teknik terimler, yontemler, malzemeler ve teknolojiler\n'
            f'- Genel kelimeler degil (analiz, sistem, gelistirme gibi) -- ozgun teknik terimler\n'
            f'- Turkce yaz, kisaltmalar buyuk harf (CFD, ROS, CFRP, LSTM gibi)\n'
            f'- Virgülle ayir, en fazla 6 kelime\n\n'
            f'YANIT (sadece virgülle ayrilmis kelimeler):'
        )
        raw   = _ext_model.generate_content(prompt).text.strip()
        lines = [l.strip() for l in raw.split('\n') if l.strip()]
        raw   = lines[-1] if lines else raw
        return [k.strip().strip('.-') for k in raw.split(',') if k.strip()][:6]
    except Exception:
        return []


def _ext_classify(title, abstract, cat_list):
    if not _EXT_OK or _ext_model is None:
        return None
    cat_str = '\n'.join(f'{i+1}. {c}' for i, c in enumerate(cat_list))
    prompt  = (
        f'Sen bir havacilk ve savunma sanayi Ar-Ge projesi siniflandirma uzmanissin.\n'
        f'Asagidaki projeyi, verilen 15 kategoriden TAM OLARAK BIRINE ata.\n\n'
        f'PROJE BASLIGI: {title}\n'
        f'PROJE OZETI: {abstract[:600]}\n\n'
        f'KATEGORILER:\n{cat_str}\n\n'
        f'KARAR KURALLARI:\n'
        f'- Projenin ana konusuna ve kullandigi yonteme odaklan\n'
        f'- Basliktaki anahtar kelimeler genellikle en guclu ipucudur\n'
        f'- Birden fazla konu varsa, en baskin olani sec\n'
        f'- Sadece listedeki kategorilerden birini sec\n\n'
        f'YANIT FORMATI (baska hicbir sey yazma):\n'
        f'SAYI. Kategori Adi\n\n'
        f'Ornek: 8. Kompozit Yapilar'
    )
    try:
        raw = _ext_model.generate_content(prompt).text.strip()
        m   = re.search(r'(\d+)\.', raw)
        if m:
            num = int(m.group(1))
            if 1 <= num <= len(cat_list):
                return cat_list[num - 1]
        raw_lower = raw.lower()
        for c in cat_list:
            if c.lower() in raw_lower:
                return c
    except Exception:
        pass
    return None


def _combine_scores(hybrid_result, extra_cat, boost=0.08):
    if not extra_cat:
        return hybrid_result
    scores = hybrid_result['all_scores']
    if extra_cat in scores:
        old           = scores[extra_cat]
        scores[extra_cat] = KategoriSkoru(
            kategori      = old.kategori,
            final_skor    = min(old.final_skor + boost, 1.0),
            keyword_skor  = old.keyword_skor,
            semantic_skor = old.semantic_skor,
            eslesmeler    = old.eslesmeler,
        )
    srt = sorted(scores.values(), key=lambda x: x.final_skor, reverse=True)
    return {
        'prediction': srt[0].kategori,
        'confidence': srt[0].final_skor,
        'top_k':      srt[:3],
        'all_scores': scores,
    }


# ── Model yukleme ────────────────────────────────────────────────────

def load_models():
    global bert_model, bert_tok, kw_model, generator, clf, TOP_KEYWORDS, ALL_CATS

    auth = {'token': HF_TOKEN} if HF_TOKEN else {}

    tax_path = hf_hub_download(
        repo_id=f'{HF_USERNAME}/liftup-bert', filename='taksonomi.txt', **auth
    )
    with open(tax_path, encoding='utf-8') as f:
        taxonomy = parse_taksonomi(f.read())
    ALL_CATS = list(taxonomy.keys())

    ckpt_path = hf_hub_download(
        repo_id=f'{HF_USERNAME}/liftup-bert', filename='checkpoint.pth', **auth
    )
    ckpt         = torch.load(ckpt_path, map_location='cpu')
    TOP_KEYWORDS = ckpt['TOP_KEYWORDS']

    bert_path = hf_hub_download(
        repo_id=f'{HF_USERNAME}/liftup-bert', filename='best_bert_model.pth', **auth
    )
    bert_tok = AutoTokenizer.from_pretrained('dbmdz/bert-base-turkish-cased')
    model    = LiftUpBertModel(len(TOP_KEYWORDS))
    model.load_state_dict(torch.load(bert_path, map_location='cpu'))
    model.eval()
    bert_model = model
    log.info('BERT hazir')

    kw_model = KeyBERT(model='sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')
    embedder = kw_model.model.embedding_model
    log.info('KeyBERT hazir')

    byt5_dir = snapshot_download(repo_id=f'{HF_USERNAME}/liftup-byt5', **auth)
    byt5_tok = AutoTokenizer.from_pretrained('google/byt5-small')
    byt5_mdl = AutoModelForSeq2SeqLM.from_pretrained(byt5_dir)
    byt5_mdl.eval()
    post     = SoftPostProcessor()

    class Generator:
        def __init__(self, tok, mdl, pp):
            self.tok = tok
            self.mdl = mdl
            self.pp  = pp

        def generate(self, title='', abstract=''):
            text = f'keywords: {title} {abstract}'.strip()
            inp  = self.tok(text, max_length=512, truncation=True, return_tensors='pt')
            with torch.no_grad():
                out = self.mdl.generate(
                    **inp, max_new_tokens=128, do_sample=False,
                    no_repeat_ngram_size=4, repetition_penalty=1.5
                )
            pred = self.tok.decode(out[0], skip_special_tokens=True)
            if pred.lower().startswith('keywords:'):
                pred = pred[9:].strip()
            kws = [k.strip() for k in pred.split(';') if k.strip()]
            return self.pp.process(kws)

    generator = Generator(byt5_tok, byt5_mdl, post)
    log.info('ByT5 hazir')

    clf = HibritSiniflandirici(
        taxonomy, embedder,
        keyword_weight=0.35, semantic_weight=0.65, title_boost=2.0
    )

    _init_ext(GEMINI_KEY, ALL_CATS)
    log.info('Tum modeller hazir')


# ── FastAPI ──────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    load_models()
    yield

app = FastAPI(title='LIFT UP Siniflandirici', lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_methods=['POST', 'GET'],
    allow_headers=['*'],
)


class ClassifyRequest(BaseModel):
    baslik:   str
    ozet:     str
    keywords: Optional[List[str]] = None


class KategoriResponse(BaseModel):
    kategori:      str
    guven:         float
    keyword_skor:  float
    semantic_skor: float
    eslesmeler:    List[str]


class ClassifyResponse(BaseModel):
    prediction:          str
    confidence:          float
    top_3:               List[KategoriResponse]
    extracted_keywords:  List[str]
    bert_keywords:       List[str]
    keybert_keywords:    List[str]
    byt5_keywords:       List[str]
    processing_time_ms:  int


def bert_extract(text):
    enc = bert_tok(
        str(text).lower(), truncation=True,
        padding='max_length', max_length=256, return_tensors='pt'
    )
    with torch.no_grad():
        logits = bert_model(enc['input_ids'], enc['attention_mask'])
        probs  = torch.sigmoid(logits)[0].numpy()
    idxs = np.argsort(probs)[-10:][::-1]
    return [TOP_KEYWORDS[i] for i in idxs if probs[i] > 0.01][:5]


def keybert_extract(text):
    clean = re.sub(r'[^\w\s\u011f\u011e\xfc\xdc\u015f\u015e\u0131\u0130\xf6\xd6\xe7\xc7]', ' ', text.lower()).strip()
    try:
        kws = kw_model.extract_keywords(
            clean, keyphrase_ngram_range=(1, 3), top_n=5, use_mmr=True, diversity=0.2
        )
        return [k[0] for k in kws][:3]
    except Exception:
        return []


@app.get('/health')
def health():
    return {'status': 'ok'}


@app.get('/')
def root():
    return {'message': 'API calisiyor', 'endpoint': 'POST /classify'}


@app.post('/classify', response_model=ClassifyResponse)
def classify(req: ClassifyRequest):
    if not req.baslik.strip() or not req.ozet.strip():
        raise HTTPException(400, 'Baslik ve ozet zorunludur')

    t0   = time.time()
    text = f'{req.baslik} {req.ozet}'

    bert_kws = bert_extract(text)
    kb_kws   = keybert_extract(text)
    byt5_kws = generator.generate(req.baslik, req.ozet)
    ext_kws  = _ext_keywords(req.baslik, req.ozet)
    extra    = req.keywords or []

    tum_kws = list(dict.fromkeys(bert_kws + kb_kws + byt5_kws + ext_kws + extra))

    hybrid  = clf.classify(tum_kws, text, title=req.baslik, top_k=3)
    ext_cat = _ext_classify(req.baslik, req.ozet, ALL_CATS)
    final   = _combine_scores(hybrid, ext_cat)

    ms = int((time.time() - t0) * 1000)

    return ClassifyResponse(
        prediction         = final['prediction'],
        confidence         = round(final['confidence'], 4),
        top_3              = [
            KategoriResponse(
                kategori      = ks.kategori,
                guven         = round(ks.final_skor, 4),
                keyword_skor  = round(ks.keyword_skor, 4),
                semantic_skor = round(ks.semantic_skor, 4),
                eslesmeler    = ks.eslesmeler[:5],
            )
            for ks in final['top_k']
        ],
        extracted_keywords = tum_kws,
        bert_keywords      = bert_kws,
        keybert_keywords   = kb_kws,
        byt5_keywords      = byt5_kws,
        processing_time_ms = ms,
    )