# src/query_builder.py

from typing import List, Dict
from langsmith import traceable


def build_queries_for_stone(stone: Dict) -> List[str]:
    """
    Tek bir taş için multiple search queries üret.
    Segmentation modeli varsa morphology ve density homogeneity de kullanılır.
    """
    size_mm = stone['size']['max_dimension_mm']
    size_category = stone['size']['size_category']
    density = stone['density']['category']

    queries = [
        # Klinik rapor bağlamı için anchor query
        f"kidney stone {size_category} {size_mm}mm clinical report findings interpretation EAU guidelines",

        # Boyut-bazlı tedavi
        f"treatment options for {size_mm}mm kidney stone clinical guidelines",

        # Boyut kategorisi yönetim
        f"management of {size_category} renal calculus recommendations",

        # Relative density karakterizasyonu (HU değil, göreli yoğunluk)
        f"{density} density kidney stone imaging characteristics composition",

        # Spontan pasaj ve MET
        f"spontaneous passage likelihood {size_mm}mm stone medical expulsive therapy",
    ]

    # Orta boy taşlar için özel sorgular (5-10mm)
    if size_category == 'medium':
        queries.append(
            f"ureteral stone {size_mm}mm spontaneous passage rate percentage probability"
        )
        queries.append(
            "kidney stone 5mm 10mm MET alpha blocker intervention threshold EAU guideline"
        )
        queries.append(
            f"kidney stone under 10mm {size_mm}mm conservative treatment observation MET option"
        )
        queries.append(
            "distal ureteral stone less than 10mm spontaneous expulsion rate alpha blocker"
        )

    # Büyük taşlar için cerrahi
    if size_category in ['large', 'very_large']:
        queries.append(
            f"surgical treatment selection {size_mm}mm renal stone EAU AUA guideline recommendation"
        )
        queries.append(
            f"PCNL indication kidney stone larger than 10mm 20mm first line EAU recommendation"
        )
        queries.append(
            f"SWL contraindication failure large renal stone size limit ESWL efficacy"
        )
        queries.append(
            "renal stone treatment algorithm 10 20mm PCNL RIRS SWL selection EAU figure"
        )

    # ── Morphology-based queries (yalnızca segmentation modeli çalıştıysa) ──
    morphology = stone.get('morphology')
    if morphology:
        circularity = morphology.get('circularity', 1.0)
        solidity = morphology.get('solidity', 1.0)

        # Düzensiz şekil → pasaj güçlüğü, obstrüksiyon riski
        if circularity < 0.65:
            queries.append(
                "irregular kidney stone morphology spontaneous passage difficulty ureteral obstruction"
            )

        # Düşük solidity → girintili-çıkıntılı yüzey
        if solidity < 0.85:
            queries.append(
                "complex spiculated kidney stone surface morphology treatment prognosis"
            )

    # ── Density homojenliği query ──
    homogeneity = stone['density'].get('homogeneity')
    if homogeneity == 'heterogeneous':
        queries.append(
            "heterogeneous kidney stone composition mixed density treatment response ESWL"
        )
    elif homogeneity == 'homogeneous':
        queries.append(
            "homogeneous kidney stone composition ESWL fragmentation efficacy"
        )

    return queries


@traceable(name="Query Builder", run_type="tool")
def build_all_queries(features: Dict) -> List[str]:
    """
    Feature extraction output'u için tüm queries.
    Unique query string listesi döndürür.
    """
    if not features.get('stone_detected', False):
        return []

    all_queries = []

    for stone in features.get('stones', []):
        all_queries.extend(build_queries_for_stone(stone))

    # Multiple stones
    if features.get('multiple_stones', False):
        all_queries.append(
            "management multiple kidney stones bilateral calculi complex treatment"
        )
        if features.get('total_stone_area_mm2'):
            all_queries.append(
                "total stone burden multiple renal calculi treatment planning"
            )

    # Duplicate'leri çıkar (case-insensitive)
    seen = set()
    unique_queries = []
    for q in all_queries:
        q_lower = q.lower()
        if q_lower not in seen:
            seen.add(q_lower)
            unique_queries.append(q)

    return unique_queries


if __name__ == "__main__":
    test_features = {
        "stone_detected": True,
        "count": 1,
        "stones": [{
            "stone_id": 1,
            "size": {
                "max_dimension_mm": 7.2,
                "size_category": "medium"
            },
            "density": {
                "category": "high",
                "homogeneity": "heterogeneous"
            },
            "shape": "irregular",
            "location": {"quadrant": "upper-left"},
            "morphology": {
                "circularity": 0.55,
                "solidity": 0.80
            }
        }],
        "multiple_stones": False
    }

    queries = build_all_queries(test_features)
    print(f"Generated {len(queries)} queries:")
    for i, q in enumerate(queries, 1):
        print(f"{i}. {q}")