import React, { useState, useEffect } from 'react';
import { TextareaControl } from "../ui/textarea-control";
import { Card, CardHeader, CardBody, CardDivider } from '../ui/wp-card';
import { Elevation } from '../ui/elevation';
import { InputControl } from '../ui/input-control';
import { Divider } from '../ui/divider';
import { DismissibleAlert } from '../ui/dismissible-alert';
import { Heading } from '../ui/heading';
import { Flex, FlexItem } from '../ui/flex';
import { Spacer } from '../ui/spacer';
import { TabPanel } from '../ui/tab-panel';
import { Disabled } from '../ui/disabled';
import { DropdownMenu, DropdownMenuItem } from '../ui/dropdown-menu';
import { useFormContext, Controller, useFieldArray } from 'react-hook-form';
import { IValueSet, useAuthServiceContext, useValueSetServiceContext } from '../../services';
import {
    FhirAge,
    FhirAnnotation,
    FhirCodeableConcept,
    FhirPeriod,
    FhirRange,
    FhirReference,
    FhirReferenceCollection,
    FhirValueSetSearch,
    TypeSelector
} from '../../components';
import {
    useFhirClinicalImpressions,
    useFhirConditionBroadcast,
    useFhirConditions,
    useFhirEncounters,
    useFhirObservations
} from '../../hooks';
import { PlusCircle, Trash2 } from 'lucide-react';
import { ChartBar } from 'lucide-react';
import { FhirResource } from 'fhir/r4';

const StageAssessmentReferences: React.FC<{ index: number; disabled: boolean; enableInlineActions: boolean; }> = ({ index, disabled, enableInlineActions = true }) => {
    const { control } = useFormContext();
    const itemReferenceFieldArray = useFieldArray({ control, name: `stage.${index}.assessment` });

    const { observationsList } = useFhirObservations();
    const { clinicalImpressionsList } = useFhirClinicalImpressions();

    return (
        <FhirReferenceCollection label="Formal record of assessment"
                                 references={itemReferenceFieldArray}
                                 context={`stage.${index}.assessment`}
                                 referenceTypes={[
                                     { name: 'ClinicalImpression', key: 'ClinicalImpression' },
                                     { name: 'Observation', key: 'Observation' }
                                 ]}
                                 getReferenceList={(referenceType: string): FhirResource[] => {
                                     switch (referenceType) {
                                         case 'ClinicalImpression':
                                             return clinicalImpressionsList;
                                         case 'Observation':
                                             return observationsList;
                                         default:
                                             return [];
                                     }
                                 }}
                                 required={false}
                                 disabled={disabled}
                                 multiple={true}
                                 enableInlineActions={enableInlineActions}/>
    );
};

interface FhirConditionProps {
    isNew: boolean;
    conditionProfile: string | undefined;
    onConditionProfileSelect: (conditionProfile: string) => void;
    elevated?: boolean;
    disabled?: boolean;
    enableInlineActions?: boolean;
}

export const FhirCondition: React.FC<FhirConditionProps> = ({ conditionProfile, onConditionProfileSelect, isNew, elevated = true, disabled = false, enableInlineActions = true }) => {
    const { control, watch } = useFormContext();
    const { fetchValueSet, valueSetUrls } = useValueSetServiceContext();
    const { userRole, USER_ROLES } = useAuthServiceContext();
    const isPatientRole = [USER_ROLES.PATIENT, USER_ROLES.RELATED_PERSON].includes(userRole as string);

    const [loading, setLoading] = useState(false);
    const [isCategoryMutable, setIsCategoryMutable] = useState<boolean>(false);
    const [categoryCodes, setCategoryCodes] = useState<IValueSet[]>([]);

    // value sets
    const [conditionClinical, setConditionClinical] = useState<IValueSet[]>([]);
    const [conditionVerStatus, setConditionVerStatus] = useState<IValueSet[]>([]);
    const [conditionCategory, setConditionCategory] = useState<IValueSet[]>([]);
    const [conditionSeverity, setConditionSeverity] = useState<IValueSet[]>([]);
    const [usCoreProblemOrHealthConcern, setUsCoreProblemOrHealthConcern] = useState<IValueSet[]>([]);
    const [usCoreSimpleObservationCategory, setUsCoreSimpleObservationCategory] = useState<IValueSet[]>([]);

    const spacerMargin = '24px';

    useFhirConditions();
    useFhirConditionBroadcast();
    const { encountersList } = useFhirEncounters();

    const onsetType = watch('onsetType');
    const abatementType = watch('abatementType');
    const stageFieldArray = useFieldArray({ control, name: 'stage' });
    const evidenceFieldArray = useFieldArray({ control, name: 'evidence' });

    useEffect(() => {
        const loadValueSets = async () => {
            setLoading(true);
            try {
                const [
                    conditionClinicalValueSet,
                    conditionVerStatusValueSet,
                    conditionSeverityValueSet,
                    conditionCategoryValueSet,
                    usCoreProblemOrHealthConcernValueSet,
                    usCoreSimpleObservationCategoryValueSet
                ] = await Promise.all([
                    fetchValueSet(valueSetUrls.conditionClinical),
                    fetchValueSet(valueSetUrls.conditionVerStatus),
                    fetchValueSet(valueSetUrls.conditionSeverity),
                    fetchValueSet(valueSetUrls.conditionCategory),
                    fetchValueSet(valueSetUrls.usCoreProblemOrHealthConcern),
                    fetchValueSet(valueSetUrls.usCoreSimpleObservationCategory)
                ]);

                setConditionClinical(conditionClinicalValueSet);
                setConditionVerStatus(conditionVerStatusValueSet);
                setConditionSeverity(conditionSeverityValueSet);
                setConditionCategory(conditionCategoryValueSet);
                setUsCoreProblemOrHealthConcern(usCoreProblemOrHealthConcernValueSet);
                setUsCoreSimpleObservationCategory(usCoreSimpleObservationCategoryValueSet);
            } catch (err) {
                console.error('Error loading value sets:', err);
            } finally {
                setLoading(false);
            }
        };

        loadValueSets();
    }, [fetchValueSet, valueSetUrls]);

    useEffect(() => {
        if (conditionProfile === 'us-core-condition-encounter-diagnosis') {
            setIsCategoryMutable(false);
            setCategoryCodes(conditionCategory);
        } else if (conditionProfile === 'us-core-condition-problems-health-concerns') {
            setIsCategoryMutable(true);
            setCategoryCodes([
                ...usCoreProblemOrHealthConcern,
                ...usCoreSimpleObservationCategory
            ])
        }
    }, [conditionProfile, setIsCategoryMutable, setIsCategoryMutable, conditionCategory, usCoreProblemOrHealthConcern, usCoreSimpleObservationCategory]);

    const addStage = () => {
        stageFieldArray.append({
            type: null,
            summary: null,
            assessment: []
        });
    };

    const removeStage = (index: number) => {
        stageFieldArray.remove(index);
    };

    const addEvidence = () => {
        evidenceFieldArray.append({
            code: []
        });
    };

    const removeEvidence = (index: number) => {
        evidenceFieldArray.remove(index);
    };

    const getSelectedConditionProfile = () => {
        switch (conditionProfile) {
            case 'us-core-condition-encounter-diagnosis':
                return 'Encounter Diagnosis';
            case 'us-core-condition-problems-health-concerns':
                return 'Problems and Health Concerns';
        }

        return '';
    };

    return (
        <Card style={elevated ? { boxShadow: 'none' } : {}}>
            {elevated && <Elevation value={5}/>}
            <CardHeader>
                <Flex>
                    <FlexItem>
                        <Heading level={5}>Details
                            {conditionProfile && (`: ${getSelectedConditionProfile()}`)}
                        </Heading>
                    </FlexItem>
                    <FlexItem>
                        {/* Radix closes the menu on select, so WP's manual onClose() is gone. */}
                        {isNew && (
                            <DropdownMenu icon={<ChartBar className="size-4" aria-hidden="true"/>}
                                          text={`${conditionProfile ? 'Change' : 'Select'} condition profile`}
                                          label={`${conditionProfile ? 'Change' : 'Select'} condition profile`}>
                                <DropdownMenuItem
                                    onSelect={() => onConditionProfileSelect('us-core-condition-encounter-diagnosis')}>
                                    Encounter Diagnosis
                                </DropdownMenuItem>
                                <DropdownMenuItem
                                    onSelect={() => onConditionProfileSelect('us-core-condition-problems-health-concerns')}>
                                    Problems and Health Concerns
                                </DropdownMenuItem>
                            </DropdownMenu>
                        )}
                    </FlexItem>
                </Flex>
            </CardHeader>
            {!conditionProfile && (
                <CardBody style={{ minHeight: 400 }}>
                    <Spacer margin={spacerMargin}/>

                    <DismissibleAlert status="info">
                        No condition profile selected.
                    </DismissibleAlert>

                    <Spacer margin={spacerMargin}/>
                </CardBody>
            )}

            {(isNew && conditionProfile || !isNew) && (
                <>
                    <CardBody>
                        {/* code */}
                        <FhirValueSetSearch label="Identification of the condition, problem or diagnosis"
                                            context="code"
                                            valueSetUrl={valueSetUrls.conditionCode}
                                            placeholder="Search for conditions..."
                                            required={true}
                                            multiple={false}
                                            disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        <Disabled isDisabled={!isCategoryMutable}>
                            {/* category */}
                            <FhirCodeableConcept label={`Category ${isCategoryMutable ? '' : '(Read Only)'}`}
                                                 context="category"
                                                 options={categoryCodes}
                                                 disabled={disabled}
                                                 loading={loading}
                                                 required={true}
                                                 multiple={true}/>
                        </Disabled>

                        <Spacer margin={spacerMargin}/>

                        {/* severity */}
                        <FhirCodeableConcept label="Subjective severity of condition"
                                             context="severity"
                                             options={conditionSeverity}
                                             loading={loading}
                                             disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        {/* clinicalStatus */}
                        <FhirCodeableConcept label="Clinical Status"
                                             context="clinicalStatus"
                                             options={conditionClinical}
                                             loading={loading}
                                             disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        {/* verificationStatus */}
                        <FhirCodeableConcept label="Verification Status"
                                             context="verificationStatus"
                                             options={conditionVerStatus}
                                             loading={loading}
                                             disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        {/* bodySite */}
                        <FhirValueSetSearch label="Anatomical location, if relevant"
                                            context="bodySite"
                                            valueSetUrl={valueSetUrls.bodySite}
                                            placeholder="Search for body sites..."
                                            multiple={true}
                                            disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        {/* encounter */}
                        {!isPatientRole && <>
                            <FhirReference label="Encounter created as part of"
                                           context="encounter"
                                           referenceType="Encounter"
                                           referenceList={encountersList}
                                           enableInlineActions={enableInlineActions}
                                           required={false}
                                           disabled={disabled}/>

                            <Spacer margin={spacerMargin}/>
                        </>}

                        {/* recordedDate */}
                        <Controller key="recordedDate"
                                    name="recordedDate"
                                    control={control}
                                    disabled={disabled}
                                    render={({ field }) => (
                                        <InputControl
                                                      {...field}
                                                      label="Date record was first recorded"
                                                      type="datetime-local"/>
                                    )}/>

                        <Spacer margin={spacerMargin}/>

                        <Heading level={5}>When did the condition begin?</Heading>
                        <Divider margin="2"/>
                        <Spacer margin={spacerMargin}/>

                        {/* onsetType */}
                        <TypeSelector label="Estimated or actual date, date-time, or age"
                                      context="onsetType"
                                      options={[
                                          { name: 'Date', key: 'dateTime' },
                                          { name: 'Age', key: 'age' },
                                          { name: 'Period', key: 'period' },
                                          { name: 'Range', key: 'range' },
                                          { name: 'String', key: 'string' }
                                      ]}
                                      required={false}
                                      disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        {onsetType === 'dateTime' && (
                            <>
                                {/* onsetDateTime */}
                                <Controller key="onsetDateTime"
                                            name="onsetDateTime"
                                            control={control}
                                            disabled={disabled}
                                            render={({ field }) => (
                                                <InputControl
                                                              {...field}
                                                              label="Onset Date"
                                                              type="datetime-local"/>
                                            )}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {onsetType === 'age' && (
                            <>
                                {/* onsetAge */}
                                <FhirAge context={`onsetAge`} disabled={disabled}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {onsetType === 'period' && (
                            <>
                                {/* onsetPeriod */}
                                <FhirPeriod label="Onset Period"
                                            context="onsetPeriod"
                                            disabled={disabled}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {onsetType === 'range' && (
                            <>
                                {/* onsetRange */}
                                <FhirRange label="Onset Range"
                                           context="onsetRange"
                                           disabled={disabled}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {onsetType === 'string' && (
                            <>
                                {/* onsetString */}
                                <Controller key="onsetString"
                                            name="onsetString"
                                            control={control}
                                            disabled={disabled}
                                            render={({ field }) => (
                                                <TextareaControl
                                                                 value={field.value}
                                                                 onChange={field.onChange}
                                                                 label="Onset Description"
                                                                 placeholder="Describe when the condition began"
                                                                 rows={3}/>
                                            )}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        <Heading level={5}>When did the condition resolve?</Heading>
                        <Divider margin="2" />
                        <Spacer margin={spacerMargin}/>

                        {/* abatementType */}
                        <TypeSelector label="When in resolution/remission"
                                      context="abatementType"
                                      options={[
                                          { name: 'Date', key: 'dateTime' },
                                          { name: 'Age', key: 'age' },
                                          { name: 'Period', key: 'period' },
                                          { name: 'Range', key: 'range' },
                                          { name: 'String', key: 'string' }
                                      ]}
                                      required={false}
                                      disabled={disabled}/>

                        <Spacer margin={spacerMargin}/>

                        {abatementType === 'dateTime' && (
                            <>
                                {/* abatementDateTime */}
                                <Controller key="abatementDateTime"
                                            name="abatementDateTime"
                                            control={control}
                                            disabled={disabled}
                                            render={({ field }) => (
                                                <InputControl
                                                              {...field}
                                                              label="Abatement Date"
                                                              type="datetime-local"/>
                                            )}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {abatementType === 'age' && (
                            <>
                                {/* abatementAge */}
                                <FhirAge context={`abatementAge`} disabled={disabled}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {abatementType === 'period' && (
                            <>
                                {/* abatementPeriod */}
                                <FhirPeriod label="Abatement Period"
                                            context="abatementPeriod"
                                            disabled={disabled}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {abatementType === 'range' && (
                            <>
                                {/* abatementRange */}
                                <FhirRange label="Abatement Range"
                                           context="abatementRange"
                                           disabled={disabled}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}

                        {abatementType === 'string' && (
                            <>
                                {/* abatementString */}
                                <Controller key="abatementString"
                                            name="abatementString"
                                            control={control}
                                            disabled={disabled}
                                            render={({ field }) => (
                                                <TextareaControl
                                                                 value={field.value}
                                                                 onChange={field.onChange}
                                                                 label="Abatement Description"
                                                                 placeholder="Describe when the condition resolved"
                                                                 rows={3}/>
                                            )}/>

                                <Spacer margin={spacerMargin}/>
                            </>
                        )}
                    </CardBody>
                    <CardDivider/>
                    <CardBody>
                        <Spacer margin={spacerMargin}/>

                        <TabPanel initialTabName="stages"
                                  tabs={[
                                      { name: 'stages', title: 'Stages' },
                                      { name: 'evidences', title: 'Evidences' },
                                      { name: 'notes', title: 'Notes' }
                                  ]}>
                            {(tab) => {
                                // stages
                                if (tab.name === 'stages') {
                                    return (
                                        <div style={{ paddingTop: '16px' }}>
                                            {stageFieldArray.fields.length === 0 && (
                                                <>
                                                    <Flex>
                                                        <FlexItem>
                                                            <Heading level={5}>Clinical stage or grade of a condition</Heading>
                                                        </FlexItem>
                                                        <FlexItem>
                                                            <button type="button"
                                                                    className={`btn btn-blue-ghost${disabled ? ' btn-busy' : ''}`}
                                                                    onClick={addStage}
                                                                    disabled={disabled}>
                                                                <PlusCircle className="w-4 h-4" />
                                                                Add Stage
                                                            </button>
                                                        </FlexItem>
                                                    </Flex>
                                                    <Divider margin="2"/>
                                                    <Spacer margin="8"/>
                                                    <DismissibleAlert status="info">
                                                        No stage added yet.
                                                    </DismissibleAlert>
                                                    <Spacer margin={spacerMargin}/>
                                                </>
                                            )}

                                            {stageFieldArray.fields.map((field, index) => (
                                                <React.Fragment key={field.id}>
                                                    <Card>
                                                        <CardHeader>
                                                            <Flex align="center" justify="space-between">
                                                                <FlexItem>
                                                                    <h4>Stage {index + 1}</h4>
                                                                </FlexItem>
                                                                <FlexItem>
                                                                    <button type="button"
                                                                            className="btn btn-red-ghost"
                                                                            title="Remove Stage"
                                                                            aria-label="Remove Stage"
                                                                            onClick={() => removeStage(index)}
                                                                            disabled={disabled}>
                                                                        <Trash2 className="w-4 h-4" />
                                                                    </button>
                                                                </FlexItem>
                                                            </Flex>
                                                        </CardHeader>
                                                        <CardBody>
                                                            <Spacer margin={spacerMargin}/>

                                                            {/* stage.type */}
                                                            <FhirValueSetSearch label="Kind of staging"
                                                                                context={`stage.${index}.type`}
                                                                                valueSetUrl={valueSetUrls.conditionStageType}
                                                                                multiple={false}
                                                                                placeholder="Search for stage types..."
                                                                                required={true}
                                                                                disabled={disabled}/>

                                                            <Spacer margin={spacerMargin}/>

                                                            {/* stage.summary */}
                                                            <FhirValueSetSearch label="Simple summary"
                                                                                context={`stage.${index}.summary`}
                                                                                valueSetUrl={valueSetUrls.conditionStage}
                                                                                multiple={false}
                                                                                placeholder="Search for stage summaries..."
                                                                                required={true}
                                                                                disabled={disabled}/>

                                                            <Spacer margin={spacerMargin}/>

                                                            {/* stage.assessment */}
                                                            <StageAssessmentReferences index={index}
                                                                                       disabled={disabled}
                                                                                       enableInlineActions={enableInlineActions}/>

                                                            <Spacer margin={spacerMargin}/>
                                                        </CardBody>
                                                    </Card>
                                                    <Spacer margin={spacerMargin} />
                                                </React.Fragment>
                                            ))}

                                            {stageFieldArray.fields.length > 0 && (
                                                <Flex justify="center">
                                                    <FlexItem>
                                                        <button type="button"
                                                                className={`btn btn-blue-ghost${disabled ? ' btn-busy' : ''}`}
                                                                onClick={addStage}
                                                                disabled={disabled}>
                                                            <PlusCircle className="w-4 h-4" />
                                                            Add Stage
                                                        </button>
                                                    </FlexItem>
                                                </Flex>
                                            )}
                                        </div>
                                    );
                                }

                                // evidences
                                if (tab.name === 'evidences') {
                                    return (
                                        <div style={{ paddingTop: '16px' }}>
                                            {evidenceFieldArray.fields.length === 0 && (
                                                <>
                                                    <Flex>
                                                        <FlexItem>
                                                            <Heading level={5}>Supporting evidence</Heading>
                                                        </FlexItem>
                                                        <FlexItem>
                                                            <button type="button"
                                                                    className={`btn btn-blue-ghost${disabled ? ' btn-busy' : ''}`}
                                                                    onClick={addEvidence}
                                                                    disabled={disabled}>
                                                                <PlusCircle className="w-4 h-4" />
                                                                Add Evidence
                                                            </button>
                                                        </FlexItem>
                                                    </Flex>
                                                    <Divider margin="2"/>
                                                    <Spacer margin="8"/>
                                                    <DismissibleAlert status="info">
                                                        No evidence added yet.
                                                    </DismissibleAlert>
                                                    <Spacer margin={spacerMargin}/>
                                                </>
                                            )}

                                            {evidenceFieldArray.fields.map((field, index) => (
                                                <React.Fragment key={field.id}>
                                                    <Card>
                                                        <CardHeader>
                                                            <Flex align="center" justify="space-between">
                                                                <FlexItem>
                                                                    <h4>Evidence {index + 1}</h4>
                                                                </FlexItem>
                                                                <FlexItem>
                                                                    <button type="button"
                                                                            className="btn btn-red-ghost"
                                                                            title="Remove Evidence"
                                                                            aria-label="Remove Evidence"
                                                                            onClick={() => removeEvidence(index)}
                                                                            disabled={disabled}>
                                                                        <Trash2 className="w-4 h-4" />
                                                                    </button>
                                                                </FlexItem>
                                                            </Flex>
                                                        </CardHeader>
                                                        <CardBody>
                                                            <Spacer margin={spacerMargin}/>

                                                            {/* evidence.code */}
                                                            <FhirValueSetSearch label="Manifestation/symptom"
                                                                                context={`evidence.${index}.code`}
                                                                                valueSetUrl={valueSetUrls.manifestationOrSymptom}
                                                                                multiple={true}
                                                                                placeholder="Search for manifestations or symptoms..."
                                                                                required={false}
                                                                                disabled={disabled}/>

                                                            <Spacer margin={spacerMargin}/>
                                                        </CardBody>
                                                    </Card>

                                                    <Spacer margin={spacerMargin}/>
                                                </React.Fragment>
                                            ))}

                                            {evidenceFieldArray.fields.length > 0 && (
                                                <Flex justify="center">
                                                    <FlexItem>
                                                        <button type="button"
                                                                className={`btn btn-blue-ghost${disabled ? ' btn-busy' : ''}`}
                                                                onClick={addEvidence}
                                                                disabled={disabled}>
                                                            <PlusCircle className="w-4 h-4" />
                                                            Add Evidence
                                                        </button>
                                                    </FlexItem>
                                                </Flex>
                                            )}
                                        </div>
                                    );
                                }

                                // notes
                                if (tab.name === 'notes') {
                                    return (
                                        <div style={{ paddingTop: '16px' }}>
                                            <Spacer margin={spacerMargin}/>

                                            {/* note */}
                                            <FhirAnnotation label="Additional information about the Condition"
                                                            context={`note`}
                                                            disabled={disabled}/>

                                            <Spacer margin={spacerMargin}/>
                                        </div>
                                    );
                                }
                            }}
                        </TabPanel>
                    </CardBody>
                </>
            )}
        </Card>
    );
};

export default FhirCondition;